{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Sum of Distances in Tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #tree #depth-first-search #graph #dynamic-programming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #树 #深度优先搜索 #图 #动态规划"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: sumOfDistancesInTree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #树中距离之和"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给定一个无向、连通的树。树中有 <code>n</code> 个标记为 <code>0...n-1</code> 的节点以及 <code>n-1</code>&nbsp;条边&nbsp;。</p>\n",
    "\n",
    "<p>给定整数 <code>n</code> 和数组&nbsp;<code>edges</code>&nbsp;，&nbsp;<code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>]</code>表示树中的节点&nbsp;<code>a<sub>i</sub></code>&nbsp;和&nbsp;<code>b<sub>i</sub></code>&nbsp;之间有一条边。</p>\n",
    "\n",
    "<p>返回长度为 <code>n</code> 的数组&nbsp;<code>answer</code>&nbsp;，其中&nbsp;<code>answer[i]</code>&nbsp;是树中第 <code>i</code> 个节点与所有其他节点之间的距离之和。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1:</strong></p>\n",
    "\n",
    "<p><img src=\"https://assets.leetcode.com/uploads/2021/07/23/lc-sumdist1.jpg\" /></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入: </strong>n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]\n",
    "<strong>输出: </strong>[8,12,6,10,10,10]\n",
    "<strong>解释: </strong>树如图所示。\n",
    "我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) \n",
    "也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2:</strong></p>\n",
    "<img alt=\"\" src=\"https://assets.leetcode.com/uploads/2021/07/23/lc-sumdist2.jpg\" />\n",
    "<pre>\n",
    "<strong>输入:</strong> n = 1, edges = []\n",
    "<strong>输出:</strong> [0]\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 3:</strong></p>\n",
    "<img alt=\"\" src=\"https://assets.leetcode.com/uploads/2021/07/23/lc-sumdist3.jpg\" />\n",
    "<pre>\n",
    "<strong>输入:</strong> n = 2, edges = [[1,0]]\n",
    "<strong>输出:</strong> [1,1]\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示:</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= n &lt;= 3 * 10<sup>4</sup></code></li>\n",
    "\t<li><code>edges.length == n - 1</code></li>\n",
    "\t<li><code>edges[i].length == 2</code></li>\n",
    "\t<li><code>0 &lt;= a<sub>i</sub>, b<sub>i</sub>&nbsp;&lt; n</code></li>\n",
    "\t<li><code>a<sub>i</sub>&nbsp;!= b<sub>i</sub></code></li>\n",
    "\t<li>给定的输入保证为有效的树</li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [sum-of-distances-in-tree](https://leetcode.cn/problems/sum-of-distances-in-tree/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [sum-of-distances-in-tree](https://leetcode.cn/problems/sum-of-distances-in-tree/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['6\\n[[0,1],[0,2],[2,3],[2,4],[2,5]]', '1\\n[]', '2\\n[[1,0]]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if n == 1: return [0]\n",
    "        ans = [-1 for _ in range(n)]\n",
    "        sons = {}\n",
    "        fathers = {0:-1}\n",
    "        for e in edges:\n",
    "            sons.setdefault(e[0], [])\n",
    "            sons[e[0]].append(e[1])\n",
    "            sons.setdefault(e[1], [])\n",
    "            sons[e[1]].append(e[0])\n",
    "        q = [0]\n",
    "        p = 0\n",
    "        # print(sons)\n",
    "        while p < len(q):\n",
    "            tmp = q[p]\n",
    "            p += 1\n",
    "            for stmp in sons[tmp]:\n",
    "                if stmp != fathers[tmp]:\n",
    "                    q.append(stmp)\n",
    "                    fathers[stmp] = tmp\n",
    "        root = 0\n",
    "        dp = {}\n",
    "        sz = {}\n",
    "        for node in q[::-1]:\n",
    "            dp[node] = 0\n",
    "            sz[node] = 0\n",
    "            for son in sons[node]:\n",
    "                if son == fathers[node]: continue\n",
    "                dp[node] += dp[son] + sz[son] + 1\n",
    "                sz[node] += sz[son] + 1\n",
    "        ncalc = 0\n",
    "        pos = [0 for _ in range(n)]\n",
    "        while ncalc < n:\n",
    "            if ans[root] == -1:\n",
    "                ans[root] = dp[root]\n",
    "                ncalc += 1\n",
    "            if pos[root] < len(sons[root]):\n",
    "                if sons[root][pos[root]] != fathers[root]:\n",
    "                    son = sons[root][pos[root]]\n",
    "                    # exchange root and son\n",
    "                    sz[root] -= sz[son] + 1\n",
    "                    dp[root] -= dp[son] + sz[son] + 1\n",
    "                    sz[son] += sz[root] + 1\n",
    "                    dp[son] += dp[root] + sz[root] + 1\n",
    "                    pos[root] += 1\n",
    "                    root = son\n",
    "                    continue\n",
    "                else:\n",
    "                    pos[root] += 1\n",
    "                    continue\n",
    "            else:\n",
    "                father = fathers[root]\n",
    "                # exchange root and father\n",
    "                sz[root] -= sz[father] + 1\n",
    "                dp[root] -= dp[father] + sz[father] + 1\n",
    "                sz[father] += sz[root] + 1\n",
    "                dp[father] += dp[root] + sz[root] + 1\n",
    "                root = father\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class Rerooting:\n",
    "    def __init__(self, n: int, decrement: int = 0):\n",
    "\n",
    "        self.n = n\n",
    "        self.adj = [[] for _ in range(n)]\n",
    "        self.root = None\n",
    "        self.decrement = decrement\n",
    "\n",
    "    def add_edge(self, u: int, v: int):\n",
    "\n",
    "        u -= self.decrement\n",
    "        v -= self.decrement\n",
    "        self.adj[u].append(v)\n",
    "        self.adj[v].append(u)\n",
    "\n",
    "    def rerooting(self, op, merge, e, root: int = 0) -> list:\n",
    "\n",
    "        root -= self.decrement\n",
    "        assert 0 <= root < self.n\n",
    "        self.root = root\n",
    "        self.parent = [-1] * self.n\n",
    "        self.order = [root]\n",
    "        stack = [root]\n",
    "        while stack:\n",
    "            from_node = stack.pop()\n",
    "            for to_node in self.adj[from_node]:\n",
    "                if to_node == self.parent[from_node]:\n",
    "                    continue\n",
    "                self.parent[to_node] = from_node\n",
    "                self.order.append(to_node)\n",
    "                stack.append(to_node)\n",
    "\n",
    "        dp1 = [e] * self.n\n",
    "        dp2 = [e] * self.n\n",
    "        for from_node in self.order[::-1]:\n",
    "            t = e\n",
    "            for to_node in self.adj[from_node]:\n",
    "                if self.parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = t\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            t = e\n",
    "            for to_node in self.adj[from_node][::-1]:\n",
    "                if self.parent[from_node] == to_node:\n",
    "                    continue\n",
    "                dp2[to_node] = merge(t, dp2[to_node])\n",
    "                t = merge(t, op(dp1[to_node], from_node, to_node))\n",
    "            dp1[from_node] = t\n",
    "\n",
    "        for new_root in self.order[1:]:\n",
    "            par = self.parent[new_root]\n",
    "            dp2[new_root] = op(merge(dp2[new_root], dp2[par]), new_root, par)\n",
    "            dp1[new_root] = merge(dp1[new_root], dp2[new_root])\n",
    "        return dp1\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        T = Rerooting(n, 0)\n",
    "        for u, v in edges:\n",
    "            T.add_edge(u, v)\n",
    "\n",
    "        def op(dp_v, u, v):\n",
    "\n",
    "            distSum, groupSize = dp_v\n",
    "            return (distSum + groupSize, groupSize + 1)\n",
    "\n",
    "        def merge(dp_pu, dp_pv):\n",
    "\n",
    "            sum1, size1 = dp_pu\n",
    "            sum2, size2 = dp_pv\n",
    "            return (sum1 + sum2, size1 + size2 - 1)\n",
    "\n",
    "        e = (0, 1)\n",
    "        dp = T.rerooting(op=op, merge=merge, e=e, root=0)\n",
    "        return [i[0] for i in dp]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if n == 1:\n",
    "            return [0]\n",
    "\n",
    "        #每个点有哪些边\n",
    "        de = defaultdict(set)\n",
    "        for e1, e2 in edges:\n",
    "            de[e1].add(e2)\n",
    "            de[e2].add(e1)\n",
    "\n",
    "        #每个点层数及下层节点个数（包括自己），从1开始\n",
    "        dc = {i: [n, 1] for i in range(n)}\n",
    "        cs = set()\n",
    "\n",
    "        for p, s in de.items():\n",
    "            if len(s) == 1:\n",
    "                cs.add(p)\n",
    "                dc[p][0] = 1\n",
    "\n",
    "        for p in cs:\n",
    "            for one in de[p]:\n",
    "                if dc[one][0] > 1:\n",
    "                    dc[one][1] += 1\n",
    "\n",
    "        now = 1\n",
    "        while cs:\n",
    "            now += 1\n",
    "            _cs = set()\n",
    "\n",
    "            for one in cs:\n",
    "                for p in de[one]:\n",
    "                    if dc[p][0] < now:\n",
    "                        continue\n",
    "\n",
    "                    s = de[p]\n",
    "                    num = 0\n",
    "                    for one in s:\n",
    "                        if dc[one][0] >= now:\n",
    "                            num += 1\n",
    "\n",
    "                    if num <= 1:\n",
    "                        _cs.add(p)\n",
    "                        dc[p][0] = now\n",
    "                        num = dc[p][0]\n",
    "\n",
    "            for p in _cs:\n",
    "                num = dc[p][1]\n",
    "                for one in de[p]:\n",
    "                    if dc[one][0] > now:\n",
    "                        dc[one][1] += num\n",
    "\n",
    "            cs = _cs\n",
    "\n",
    "        #每条边两端的点数\n",
    "        def get(p1, p2):\n",
    "            c1, num1 = dc[p1]\n",
    "            c2, num2 = dc[p2]\n",
    "\n",
    "            if c1 <= c2:\n",
    "                return num1\n",
    "            return n - num2\n",
    "\n",
    "        #计算点0\n",
    "        d0 = {0: 0}\n",
    "        done = {0}\n",
    "        jl = 0\n",
    "\n",
    "        while done:\n",
    "            now = set()\n",
    "            jl += 1\n",
    "\n",
    "            for p in done:\n",
    "                for one in de[p]:\n",
    "                    if one not in d0:\n",
    "                        d0[one] = jl\n",
    "                        now.add(one)\n",
    "\n",
    "            done = now\n",
    "\n",
    "        ret = [0]*n\n",
    "        ret[0] = sum(d0.values())\n",
    "\n",
    "        done = {0}\n",
    "        while done:\n",
    "            now = set()\n",
    "\n",
    "            for p in done:\n",
    "                retp = ret[p]\n",
    "\n",
    "                for one in de[p]:\n",
    "                    if ret[one] != 0:\n",
    "                        continue\n",
    "\n",
    "                    num = get(p, one)\n",
    "                    ret[one] = retp + num - n + num\n",
    "                    now.add(one)\n",
    "\n",
    "            done = now\n",
    "\n",
    "        return ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]], childd=None, childc=None ,rn=None) -> List[int]:\n",
    "        if n==1: return [0]\n",
    "        if len(edges)==1:\n",
    "            a,b = edges[0]\n",
    "            if childd is None:\n",
    "                return [1,1]\n",
    "            s = childd[a] + childd[b]\n",
    "            childd[a],childd[b] = s+childc[b]+1,s+childc[a]+1\n",
    "            return childd\n",
    "        d = [0]*n\n",
    "        neigh = [[] for _ in range(n)]\n",
    "        for i,j in edges:\n",
    "            d[i]+=1\n",
    "            d[j]+=1\n",
    "            neigh[i].append(j)\n",
    "            neigh[j].append(i)\n",
    "        f = [i for i in range(n) if d[i]==1]\n",
    "        v = set(f)\n",
    "        n2l = [0]*n\n",
    "        l2n = defaultdict(list)\n",
    "        parent = dict()\n",
    "        dc = defaultdict(int)\n",
    "        if childc is None: childc = [0]*n\n",
    "        if childd is None: childd = [0]*n\n",
    "        tr = [[i for i in f]]\n",
    "        while len(v)<n:\n",
    "            nex = set()\n",
    "            for i in f:\n",
    "                for j in neigh[i]:\n",
    "                    if j not in v and len(neigh[j])-dc[j]==2:\n",
    "                        # higher level\n",
    "                        nex.add(j)\n",
    "                        v.add(j)\n",
    "                    if j in nex or j not in v:\n",
    "                        dc[j]+=1\n",
    "                        parent[i] = j\n",
    "                        childc[j]+= childc[i]+1\n",
    "                        childd[j]+=childd[i]+childc[i]+1 # i's dis\n",
    "                n2l[i] = len(tr)-1\n",
    "            tr.append(list(nex))\n",
    "            if len(v)==n:\n",
    "                break\n",
    "            f = nex\n",
    "        if len(tr[-1])>1:\n",
    "            a,b = tr[-1]\n",
    "            tr.append([tr[-1].pop()])\n",
    "            parent[a] = b\n",
    "            childc[b] += childc[a] + 1\n",
    "            childd[b] += childd[a] + childc[a] + 1 \n",
    "        if rn is None: rn = n\n",
    "        for i in range(len(tr)-2,-1,-1):\n",
    "            for j in tr[i]:\n",
    "                p = parent[j]\n",
    "                childd[j] += (childd[p]-childd[j]-childc[j])+rn-childc[j]-2\n",
    "        return childd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import numpy as np\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:\n",
    "        # 树形dp\n",
    "        if N == 0:\n",
    "            return None\n",
    "        if N == 1:\n",
    "            return [0]\n",
    "        self.dp = [0] * N   # dp[i]为当前结构下节点i到其他节点距离\n",
    "        self.sz = [0] * N   # sz[i]为当前子树的大小\n",
    "        self.res = [0] * N  \n",
    "    \n",
    "        rela = {}\n",
    "        for edge in edges:\n",
    "            if edge[0] not in rela:\n",
    "                rela[edge[0]] = [edge[1]]\n",
    "            else:\n",
    "                rela[edge[0]].append(edge[1])\n",
    "            if edge[1] not in rela:\n",
    "                rela[edge[1]] = [edge[0]]\n",
    "            else:\n",
    "                rela[edge[1]].append(edge[0])\n",
    "        \n",
    "        def dfs(n, p):\n",
    "            self.sz[n] = 1\n",
    "            self.dp[n] = 0\n",
    "            son_list = rela[n]\n",
    "            for son in son_list:\n",
    "                if son == p:\n",
    "                    continue\n",
    "                dfs(son, n)\n",
    "                self.sz[n] += self.sz[son]\n",
    "                self.dp[n] += self.dp[son] + self.sz[son]\n",
    "            return\n",
    "\n",
    "        dfs(0, -1)\n",
    "\n",
    "        def dfs2(n, p):\n",
    "            self.res[n] = self.dp[n]   # 记录n为root时的情况，接下来将n的直接子节点转换为root\n",
    "            son_list = rela[n]\n",
    "            for son in son_list:\n",
    "                if son == p:    # p为当前情况之前考虑过的情况，即之前p是n的父节点\n",
    "                    continue   \n",
    "                # 将son转为root\n",
    "                sz_n, dp_n = self.sz[n], self.dp[n]\n",
    "                sz_son, dp_son = self.sz[son], self.dp[son]\n",
    "                self.sz[n] -= self.sz[son]\n",
    "                self.dp[n] -= self.dp[son] + self.sz[son]\n",
    "                self.sz[son] += self.sz[n]\n",
    "                self.dp[son] += self.dp[n] + self.sz[n]\n",
    "                # 转换完成后dfs\n",
    "                dfs2(son, n)\n",
    "                # son这一分支dfs遍历完后，sz，dp恢复到原状\n",
    "                self.dp[n], self.dp[son], self.sz[n], self.sz[son] = dp_n, dp_son, sz_n, sz_son\n",
    "            return\n",
    "        \n",
    "        dfs2(0, -1)\n",
    "        return self.res\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        # ------------------------------------\n",
    "        # 方法一：floyd  O(N^3)超时\n",
    "        INF = np.inf\n",
    "        dis = [[INF] * N for _ in range(N)]   # dis[i][j]为当前找到的i到j的最短距离\n",
    "        for i in range(N):\n",
    "            dis[i][i] = 0\n",
    "        for edge in edges:\n",
    "            dis[edge[0]][edge[1]] = 1\n",
    "            dis[edge[1]][edge[0]] = 1\n",
    "        for k in range(N):\n",
    "            for i in range(N):\n",
    "                for j in range(N):\n",
    "                    if dis[i][j] > dis[i][k] + dis[k][j]:\n",
    "                        dis[i][j] = dis[i][k] + dis[k][j]\n",
    "        res = [sum(li) for li in dis]\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if n==1:\n",
    "            return [0]\n",
    "        con = {}\n",
    "        cnt = {}\n",
    "        accu = {}\n",
    "        for x, y in edges:\n",
    "            if x not in con:\n",
    "                con[x] = set()\n",
    "                cnt[x] = 0\n",
    "            if y not in con:\n",
    "                con[y] = set()\n",
    "                cnt[y] = 0\n",
    "            con[x].add(y)\n",
    "            cnt[x] += 1\n",
    "            con[y].add(x)\n",
    "            cnt[y] += 1\n",
    "        d = collections.deque()\n",
    "        vi = set()\n",
    "        for i in range(n):\n",
    "            accu[i] = 0\n",
    "            if cnt[i] == 1:\n",
    "                d.append(i)\n",
    "                vi.add(i)\n",
    "\n",
    "\n",
    "        while d:\n",
    "            tmp = d.popleft()\n",
    "            cnt[tmp] = 0\n",
    "            vi.add(tmp)\n",
    "            for to in con[tmp]:\n",
    "                if to not in vi:\n",
    "                    accu[to] += accu[tmp] + 1\n",
    "                    cnt[to] -= 1\n",
    "                    if cnt[to] == 1:\n",
    "                        d.append(to)\n",
    "        st = collections.deque()\n",
    "        ans = [0 for _ in range(n)]\n",
    "        fi = [0] * n\n",
    "\n",
    "        st.append(0)\n",
    "        vi = set()\n",
    "        vi.add(0)\n",
    "        while st:\n",
    "            t = st.popleft()\n",
    "            for x in con[t]:\n",
    "                if x not in vi:\n",
    "                    vi.add(x)\n",
    "                    fi[x] = fi[t] + 1\n",
    "                    st.append(x)\n",
    "        ans[0] = sum(fi)\n",
    "\n",
    "        st.append(0)\n",
    "        vi = set()\n",
    "        vi.add(0)\n",
    "        while st:\n",
    "            t = st.popleft()\n",
    "            for x in con[t]:\n",
    "                if x not in vi:\n",
    "                    vi.add(x)\n",
    "                    st.append(x)\n",
    "                    if accu[x] > accu[t]:\n",
    "                        ans[x] = ans[t] - (n - 2 * (accu[t] + 1))\n",
    "                    else:\n",
    "                        ans[x] = ans[t] + (n - 2 * (accu[x] + 1))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:\n",
    "        import numpy as np\n",
    "        self.child = {}\n",
    "        for a, b in edges:\n",
    "            self.child[a] = self.child.get(a, []) + [b]\n",
    "            self.child[b] = self.child.get(b, []) + [a]\n",
    "        self.tree = dict([(i, {}) for i in range(N)])   # 存放每个节点的后代数\n",
    "        self.dis = dict([(i, {}) for i in range(N)])    # 存放每个节点的与子节点的总距离\n",
    "        self.search(0, 0)\n",
    "        dp = [0 for i in range(N)]\n",
    "        dp[0] = sum(self.dis[0].values())\n",
    "        history = {0}\n",
    "        stack = [0]\n",
    "        \n",
    "        while stack:\n",
    "            root = stack.pop(0)\n",
    "            # print(\"-----------\")\n",
    "            # print(root)\n",
    "            # print(self.tree)\n",
    "            # print(self.dis)\n",
    "            for c in self.child.get(root, []):\n",
    "                if c not in history:\n",
    "                    dp[c] = sum(self.dis[c].values()) + dp[root] - self.dis[root][c] + N - self.tree[root][c]\n",
    "                    self.tree[c][root] = N - self.tree[root][c]\n",
    "                    self.dis[c][root] = dp[root] - self.dis[root][c] + N - self.tree[root][c]\n",
    "                    history.add(c)\n",
    "                    stack.append(c)\n",
    "        return dp\n",
    "\n",
    "    def search(self, root, father, depth=1):\n",
    "        total_nodes = 0\n",
    "        total_dis = 0\n",
    "        for c in self.child.get(root, []):\n",
    "            if c != father:\n",
    "                temp = self.search(c, root)\n",
    "                total_nodes += temp[0]\n",
    "                total_dis += temp[1]\n",
    "                self.tree[root][c] = temp[0]\n",
    "                self.dis[root][c] = temp[1]\n",
    "        if total_nodes == 0:\n",
    "            return 1, 1\n",
    "        total_nodes += 1\n",
    "        return total_nodes, total_dis + total_nodes        \n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        parent_children = {}\n",
    "        for e in edges:\n",
    "            children = parent_children.get(e[0], [])\n",
    "            children.append(e[1])\n",
    "            parent_children[e[0]] = children\n",
    "\n",
    "            children = parent_children.get(e[1], [])\n",
    "            children.append(e[0])\n",
    "            parent_children[e[1]] = children\n",
    "\n",
    "        ready_count = [0 for i in range(n)]\n",
    "        distance = {}\n",
    "\n",
    "        queue = []\n",
    "        for parent in parent_children:\n",
    "            if len(parent_children[parent]) == 1:\n",
    "                queue.append(parent)\n",
    "\n",
    "        res = [0 for i in range(n)]\n",
    "\n",
    "        start = 0\n",
    "        while start < len(queue):\n",
    "            node_id = queue[start]\n",
    "            start += 1\n",
    "            children = parent_children[node_id]\n",
    "            sum_dist = 0\n",
    "            sum_child = 0\n",
    "            miss_child = None\n",
    "            for child in children:\n",
    "                if child * n + node_id in distance:\n",
    "                    [dist, n_child] = distance.get(child * n + node_id)\n",
    "                    sum_dist += dist\n",
    "                    sum_child += n_child\n",
    "                else:\n",
    "                    miss_child = child\n",
    "\n",
    "            if ready_count[node_id] == len(children) - 1:\n",
    "                sum_dist += sum_child + 1\n",
    "                distance[node_id * n + miss_child] = [sum_dist, sum_child + 1]\n",
    "                ready_count[miss_child] += 1\n",
    "                if ready_count[miss_child] >= len(parent_children[miss_child]) - 1:\n",
    "                    queue.append(miss_child)\n",
    "            else:\n",
    "                #self.show(n, distance)\n",
    "                res[node_id] = sum_dist\n",
    "                for child in children:\n",
    "                    if res[child] == 0:\n",
    "                        [dist, n_child] = distance[child * n + node_id]\n",
    "                        distance[node_id * n + child] = [sum_dist - dist + sum_child - n_child + 1, sum_child - n_child + 1]\n",
    "                        ready_count[child] += 1\n",
    "                        if ready_count[child] >= len(parent_children[child]) - 1:\n",
    "                            queue.append(child)\n",
    "\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        def default_val(i,first = True):\n",
    "            return [0 ,first ]\n",
    "        def reduce_op(res1,res2):\n",
    "            return [res1[0]+res2[0],res1[1]+res2[1]]\n",
    "        def composition(u,v,res_v) : \n",
    "            return [res_v[0] + res_v[1],res_v[1]]\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u) \n",
    "        parent = [-1] * n\n",
    "        order,stack = [0], [0]\n",
    "        while stack:\n",
    "            u = stack.pop()\n",
    "            g[u] = [v for v in g[u] if v != parent[u]] #去掉父亲\n",
    "            for v in g[u]:\n",
    "                parent[v] = u\n",
    "                order.append(v)\n",
    "                stack.append(v)  \n",
    "        down = [None] * n \n",
    "        parent_down_without_me = [None] * n \n",
    "        for me in reversed(order):\n",
    "            res = default_val(me,True)\n",
    "            for son in g[me]: \n",
    "                parent_down_without_me[son] = res \n",
    "                res = reduce_op(composition(me,son,down[son]),res)\n",
    "            down[me] = res \n",
    "            res = default_val(me,False)\n",
    "            for son in reversed(g[me]): \n",
    "                parent_down_without_me[son] = reduce_op(res,parent_down_without_me[son]) \n",
    "                res = reduce_op(composition(me,son,down[son]),res) \n",
    "        up = [default_val(0,False) ] + [None] * (n-1)\n",
    "        for i in order[1:]:\n",
    "            up[i] = composition(i,parent[i],reduce_op(up[parent[i]],parent_down_without_me[i]))\n",
    "        updown = [reduce_op(x,y) for x,y in zip(up,down)]\n",
    "        # print(up,down,parent_down_without_me,updown,sep='\\n')\n",
    "        return [x[0] for x in updown]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "\n",
    "        if n == 1:\n",
    "            return [0]\n",
    "        \n",
    "        d = {}\n",
    "        for a,b in edges:\n",
    "            if a not in d:\n",
    "                d[a] = []\n",
    "            if b not in d:\n",
    "                d[b] = []\n",
    "            d[a].append(b)\n",
    "            d[b].append(a)\n",
    "        \n",
    "        de = {}\n",
    "        parent = [k for k in d.keys() if len(d[k]) == 1]\n",
    "        visited = set([])\n",
    "\n",
    "        pool = [[0,1] for i in range(n)]\n",
    "\n",
    "        while parent:\n",
    "            child = []\n",
    "            for p in parent:\n",
    "                visited.add(p)\n",
    "                for nxt in d[p]:\n",
    "                    if nxt not in visited:\n",
    "                        de[(p,nxt)] = pool[p][1]\n",
    "                        de[(nxt,p)] = n-pool[p][1]\n",
    "                        pool[nxt][0] += 1\n",
    "                        pool[nxt][1] += pool[p][1]\n",
    "                        if pool[nxt][0] == len(d[nxt]) - 1:\n",
    "                            child.append(nxt)\n",
    "            \n",
    "            parent = child\n",
    "        \n",
    "        #print(pool)\n",
    "        #print(de)\n",
    "\n",
    "        start = [k for k in d.keys() if len(d[k]) == 1][0]\n",
    "\n",
    "        res = [0] * n\n",
    "\n",
    "        sumd = 0\n",
    "        td = 0\n",
    "        parent = [start]\n",
    "        visited = set([start])\n",
    "        while parent:\n",
    "            td += 1\n",
    "            child = []\n",
    "            for p in parent:\n",
    "                visited.add(p)\n",
    "                for nxt in d[p]:\n",
    "                    if nxt not in visited:\n",
    "                        sumd += td\n",
    "                        res[nxt] = res[p] + de[(p,nxt)] - de[(nxt,p)]\n",
    "                        child.append(nxt)\n",
    "            \n",
    "            parent = child\n",
    "        \n",
    "        for i in range(n):\n",
    "            res[i] += sumd\n",
    "\n",
    "        return res\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "\n",
    "        if n == 1:\n",
    "            return [0]\n",
    "        \n",
    "        d = {}\n",
    "        for a,b in edges:\n",
    "            if a not in d:\n",
    "                d[a] = []\n",
    "            if b not in d:\n",
    "                d[b] = []\n",
    "            d[a].append(b)\n",
    "            d[b].append(a)\n",
    "        \n",
    "        de = {}\n",
    "        parent = [k for k in d.keys() if len(d[k]) == 1]\n",
    "        visited = set([])\n",
    "\n",
    "        pool = [[0,1] for i in range(n)]\n",
    "\n",
    "        while parent:\n",
    "            child = []\n",
    "            for p in parent:\n",
    "                visited.add(p)\n",
    "                for nxt in d[p]:\n",
    "                    if nxt not in visited:\n",
    "                        de[(p,nxt)] = pool[p][1]\n",
    "                        de[(nxt,p)] = n-pool[p][1]\n",
    "                        pool[nxt][0] += 1\n",
    "                        pool[nxt][1] += pool[p][1]\n",
    "                        if pool[nxt][0] == len(d[nxt]) - 1:\n",
    "                            child.append(nxt)\n",
    "            \n",
    "            parent = child\n",
    "        \n",
    "        #print(pool)\n",
    "        #print(de)\n",
    "\n",
    "        start = [k for k in d.keys() if len(d[k]) == 1][0]\n",
    "\n",
    "        sumd = 0\n",
    "        td = 0\n",
    "        parent = [start]\n",
    "        visited = set([start])\n",
    "        while parent:\n",
    "            td += 1\n",
    "            child = []\n",
    "            for p in parent:\n",
    "                visited.add(p)\n",
    "                for nxt in d[p]:\n",
    "                    if nxt not in visited:\n",
    "                        sumd += td\n",
    "                        child.append(nxt)\n",
    "            \n",
    "            parent = child\n",
    "        \n",
    "        #print(start,sumd)\n",
    "\n",
    "        res = [0] * n\n",
    "        res[start] = sumd\n",
    "\n",
    "        parent = [[start,sumd]]\n",
    "        visited = set([start])\n",
    "        while parent:\n",
    "            child = []\n",
    "            for plast, psumd in parent:\n",
    "                visited.add(plast)\n",
    "                for nxt in d[plast]:\n",
    "                    if nxt not in visited:\n",
    "                        res[nxt] = psumd+de[(plast,nxt)]-de[(nxt,plast)]\n",
    "                        child.append([nxt,res[nxt]])\n",
    "            \n",
    "            parent = child\n",
    "        \n",
    "        return res\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if n == 1:\n",
    "            return [0]\n",
    "        \n",
    "        # 每个点有哪些边\n",
    "        le = [set() for _ in range(n)]\n",
    "        for e1, e2 in edges:\n",
    "            le[e1].add(e2)\n",
    "            le[e2].add(e1)\n",
    "\n",
    "        # 每条边一边的点数\n",
    "        dnum = {}\n",
    "        lnum = [1] * n\n",
    "        dp = [i for i, v in enumerate(le) if len(v) == 1]\n",
    "\n",
    "        # 一个边缘点\n",
    "        f = dp[0]\n",
    "\n",
    "        # 下面处理要修改le，复制一个\n",
    "        le_c = [v.copy() for v in le]\n",
    "\n",
    "        while 1:\n",
    "            now = []\n",
    "            for i in dp:\n",
    "                if not le_c[i]:\n",
    "                    continue\n",
    "\n",
    "                j = le_c[i].pop()\n",
    "                dnum[j, i] = lnum[i]\n",
    "                lnum[j] += lnum[i]\n",
    "\n",
    "                le_c[j].remove(i)\n",
    "                if len(le_c[j]) == 1:\n",
    "                    now.append(j)\n",
    "\n",
    "            if now:\n",
    "                dp = now\n",
    "            else:\n",
    "                break\n",
    "\n",
    "        # 计算到f的距离\n",
    "        jl = [-1] * n\n",
    "        dp = [f]\n",
    "        jl[f] = 0\n",
    "\n",
    "        while 1:\n",
    "            now = []\n",
    "            for i in dp:\n",
    "                for j in le[i]:\n",
    "                    if jl[j] == -1:\n",
    "                        jl[j] = jl[i] + 1\n",
    "                        now.append(j)\n",
    "\n",
    "            if now:\n",
    "                dp = now\n",
    "            else:\n",
    "                break\n",
    "\n",
    "        # 返回值\n",
    "        ret = [0] * n\n",
    "        ret[f] = sum(jl)\n",
    "        dp = [f]\n",
    "\n",
    "        while 1:\n",
    "            now = []\n",
    "            for i in dp:\n",
    "                for j in le[i]:\n",
    "                    if ret[j]:\n",
    "                        continue\n",
    "\n",
    "                    if (j, i) in dnum:\n",
    "                        inum = dnum[j, i]\n",
    "                        jnum = n - inum\n",
    "                    else:\n",
    "                        jnum = dnum[i, j]\n",
    "                        inum = n - jnum\n",
    "\n",
    "                    ret[j] = ret[i] + inum - jnum\n",
    "                    now.append(j)\n",
    "\n",
    "            if now:\n",
    "                dp = now\n",
    "            else:\n",
    "                break\n",
    "\n",
    "        return ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        ma = [[] for _ in range(n)]\n",
    "        sumdis = [0] * n\n",
    "        ans = [0] * n\n",
    "        ns = [1] * n\n",
    "        for f, t in edges:\n",
    "            ma[f].append(t)\n",
    "            ma[t].append(f)\n",
    "\n",
    "        def dfs1(f, pre):\n",
    "            for t in ma[f]:\n",
    "                if t != pre:\n",
    "                    dfs1(t, f)\n",
    "                    ns[f] += ns[t]\n",
    "                    sumdis[f] += sumdis[t] + ns[t]\n",
    "\n",
    "        def dfs2(f, pre):\n",
    "            if pre != -1:\n",
    "                ans[f] = ans[pre] + n - 2*ns[f]\n",
    "            else:\n",
    "                ans[f] = sumdis[f]\n",
    "            for t in ma[f]:\n",
    "                if t != pre:\n",
    "                    dfs2(t, f)\n",
    "\n",
    "        dfs1(n>>1, -1)\n",
    "        dfs2(n>>1, -1)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import deque\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if n == 1:\n",
    "            return [0]\n",
    "        debug = 0\n",
    "        def _ei(x,y):\n",
    "            return x*n+y\n",
    "        def _uei(x,y):\n",
    "            if x>y:\n",
    "                return _ei(y,x)\n",
    "            return _ei(x,y)\n",
    "        neis = {i:[] for i in range(n)}\n",
    "        for x,y in edges:\n",
    "            neis[x].append(y)\n",
    "            neis[y].append(x)\n",
    "        start = []\n",
    "        degrees = [0] * n\n",
    "        for i in range(n):\n",
    "            degrees[i] = len(neis[i])\n",
    "            if degrees[i] == 1:\n",
    "                start.append(i)\n",
    "        pts0 = {}\n",
    "        pts1 = {}\n",
    "        for end in start:\n",
    "            x,y = end,neis[end][0]\n",
    "            z,f = _ei(x,y),_ei(y,x)\n",
    "            pts0[z] = n-1\n",
    "            pts0[f] = 1\n",
    "            pts1[f] = 1\n",
    "        q = deque(start)\n",
    "        tempd = degrees.copy()\n",
    "        used = set()\n",
    "        idx = 0\n",
    "        orders = [-1] * n\n",
    "        edgesorder = []\n",
    "        while q:\n",
    "            now = q.popleft()\n",
    "            orders[now] = idx\n",
    "            idx += 1\n",
    "            # if debug:\n",
    "            #     print(now,q,pts0,pts1,orders)\n",
    "            if degrees[now]>1:\n",
    "                unused = maxidx = -1\n",
    "                cntnode = cntpts = 0\n",
    "                for nei in neis[now]:\n",
    "                    if nei in used:\n",
    "                        if maxidx == -1 or (maxidx != -1 and orders[nei] > orders[maxidx]):\n",
    "                            maxidx = nei\n",
    "                        cntnode += pts0[_ei(now,nei)]\n",
    "                        cntpts += pts1[_ei(now,nei)]\n",
    "                    else:\n",
    "                        unused = nei\n",
    "                \n",
    "                if unused == -1:\n",
    "                    unused = maxidx\n",
    "                    z,f = _ei(now,unused),_ei(unused,now)\n",
    "                    cntnode -= pts0[z]\n",
    "                    cntpts -= pts1[z]\n",
    "                    pts1[f] = cntnode+cntpts+1\n",
    "                else:\n",
    "                    z,f = _ei(now,unused),_ei(unused,now)\n",
    "                    pts0[f] = cntnode+1\n",
    "                    pts0[z] = n-cntnode-1\n",
    "                    pts1[f] = cntnode+cntpts+1\n",
    "                # if debug:\n",
    "                #     print(now,unused,used,neis[now],cntnode,cntpts)\n",
    "            tempd[now] = 0\n",
    "            used.add(now)\n",
    "            for nei in neis[now]:\n",
    "                if tempd[nei]:\n",
    "                    tempd[nei] -= 1\n",
    "                    edgesorder.append(_ei(now,nei))\n",
    "                    if tempd[nei] == 1:\n",
    "                        q.append(nei)\n",
    "        # return 0\n",
    "        # if debug:\n",
    "        #     print(len(edgesorder))\n",
    "            # return 0\n",
    "        out = [-1]*n\n",
    "        while edgesorder:\n",
    "            edgekey = edgesorder.pop()\n",
    "            for now in [edgekey//n,edgekey%n]:\n",
    "                if out[now] != -1:\n",
    "                    continue\n",
    "                out[now] = 0\n",
    "                for nei in neis[now]:\n",
    "                    out[now] += pts1[_ei(now,nei)]\n",
    "                for nei in neis[now]:\n",
    "                    k = _ei(nei,now)\n",
    "                    if k not in pts1:\n",
    "                        pts1[k] = out[now] - pts1[_ei(now,nei)] + pts0[_ei(nei,now)]\n",
    "        # if debug:\n",
    "        #     print([[k//n,k%n,v] for k,v in pts0.items()])\n",
    "        #     print([[k//n,k%n,v] for k,v in pts1.items()])\n",
    "        return out\n",
    "'''\n",
    "6\n",
    "[[0,1],[0,2],[2,3],[2,4],[2,5]]\n",
    "1\n",
    "[]\n",
    "2\n",
    "[[1,0]]\n",
    "'''\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import deque\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if n == 1:\n",
    "            return [0]\n",
    "        debug = 0\n",
    "        def _ei(x,y):\n",
    "            return x*n+y\n",
    "        neis = {i:[] for i in range(n)}\n",
    "        for x,y in edges:\n",
    "            neis[x].append(y)\n",
    "            neis[y].append(x)\n",
    "        start = []\n",
    "        degrees = [0] * n\n",
    "        for i in range(n):\n",
    "            degrees[i] = len(neis[i])\n",
    "            if degrees[i] == 1:\n",
    "                start.append(i)\n",
    "        pts0 = {}\n",
    "        pts1 = {}\n",
    "        for end in start:\n",
    "            x,y = end,neis[end][0]\n",
    "            z,f = _ei(x,y),_ei(y,x)\n",
    "            pts0[z] = n-1\n",
    "            pts0[f] = 1\n",
    "            pts1[f] = 1\n",
    "        q = deque(start)\n",
    "        tempd = degrees.copy()\n",
    "        used = set()\n",
    "        idx = 0\n",
    "        orders = [-1] * n\n",
    "        edgesorder = []\n",
    "        while q:\n",
    "            now = q.popleft()\n",
    "            orders[now] = idx\n",
    "            idx += 1\n",
    "            # if debug:\n",
    "            #     print(now,q,pts0,pts1,orders)\n",
    "            if degrees[now]>1:\n",
    "                unused = maxidx = -1\n",
    "                cntnode = cntpts = 0\n",
    "                for nei in neis[now]:\n",
    "                    if nei in used:\n",
    "                        if maxidx == -1 or (maxidx != -1 and orders[nei] > orders[maxidx]):\n",
    "                            maxidx = nei\n",
    "                        cntnode += pts0[_ei(now,nei)]\n",
    "                        cntpts += pts1[_ei(now,nei)]\n",
    "                    else:\n",
    "                        unused = nei\n",
    "                \n",
    "                if unused == -1:\n",
    "                    unused = maxidx\n",
    "                    z,f = _ei(now,unused),_ei(unused,now)\n",
    "                    cntnode -= pts0[z]\n",
    "                    cntpts -= pts1[z]\n",
    "                    pts1[f] = cntnode+cntpts+1\n",
    "                else:\n",
    "                    z,f = _ei(now,unused),_ei(unused,now)\n",
    "                    pts0[f] = cntnode+1\n",
    "                    pts0[z] = n-cntnode-1\n",
    "                    pts1[f] = cntnode+cntpts+1\n",
    "            tempd[now] = 0\n",
    "            used.add(now)\n",
    "            for nei in neis[now]:\n",
    "                if tempd[nei]:\n",
    "                    tempd[nei] -= 1\n",
    "                    edgesorder.append(_ei(now,nei))\n",
    "                    if tempd[nei] == 1:\n",
    "                        q.append(nei)\n",
    "        out = [-1]*n\n",
    "        while edgesorder:\n",
    "            edgekey = edgesorder.pop()\n",
    "            for now in [edgekey//n,edgekey%n]:\n",
    "                if out[now] != -1:\n",
    "                    continue\n",
    "                out[now] = 0\n",
    "                for nei in neis[now]:\n",
    "                    out[now] += pts1[_ei(now,nei)]\n",
    "                for nei in neis[now]:\n",
    "                    k = _ei(nei,now)\n",
    "                    if k not in pts1:\n",
    "                        pts1[k] = out[now] - pts1[_ei(now,nei)] + pts0[_ei(nei,now)]\n",
    "        return out\n",
    "'''\n",
    "6\n",
    "[[0,1],[0,2],[2,3],[2,4],[2,5]]\n",
    "1\n",
    "[]\n",
    "2\n",
    "[[1,0]]\n",
    "'''\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def dfs_forward(self, point, root, forwards, links):\n",
    "        # forwards[point] = [1, 0]\n",
    "        forwards[point] = 1\n",
    "        distance = 0\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                num, dis = self.dfs_forward(p, point, forwards, links)\n",
    "                # forwards[point][0] += forwards[p][0] # num\n",
    "                # forwards[point][1] += (forwards[p][0] + forwards[p][1]) # distance_sum\n",
    "                forwards[point] += num\n",
    "                distance += num + dis\n",
    "        return forwards[point], distance\n",
    "\n",
    "    def dfs_backward(self, point, root, backwards, forwards, links):\n",
    "        # + (n-k) - k \n",
    "        backwards[point] = backwards[root] + len(backwards) - 2 * forwards[point]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_backward(p, point, backwards, forwards, links)\n",
    "        return  \n",
    "\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if len(edges) == 0:\n",
    "            return [0]\n",
    "        # links = collections.defaultdict(list)\n",
    "        links = [[] for _ in range(n)]\n",
    "        for p1, p2 in edges:\n",
    "            links[p1].append(p2)\n",
    "            links[p2].append(p1)\n",
    "        p0 = edges[0][0]\n",
    "        forwards = [0 for i in range(n)]\n",
    "        backwards = [0 for i in range(n)]\n",
    "        n_total, backwards[p0] = self.dfs_forward(p0, -1, forwards, links)\n",
    "\n",
    "        for p in links[p0]:\n",
    "            self.dfs_backward(p, p0, backwards, forwards, links)\n",
    "        return backwards\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        g = [[] for _ in range(n)]  # g[x] 表示 x 的所有邻居\n",
    "        \n",
    "\n",
    "        # k: random root.\n",
    "        root = random.randint(0, n-1)\n",
    "\n",
    "        # 邻接表\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "\n",
    "        # for i in range(n):\n",
    "        #     if len(g[i]) > len(g[k]):\n",
    "        #         k = i\n",
    "\n",
    "        ans = [0] * n\n",
    "        size = [1] * n  # 注意这里初始化成 1 了，下面只需要累加儿子的子树大小\n",
    "        def dfs(x: int, fa: int, depth: int) -> None:\n",
    "            ans[root] += depth  # depth 为 k 到 x 的距离\n",
    "            for y in g[x]:  # 遍历 x 的邻居 y\n",
    "                if y != fa:  # 避免访问父节点\n",
    "                    dfs(y, x, depth + 1)  # x 是 y 的父节点\n",
    "                    size[x] += size[y]  # 累加 x 的儿子 y 的子树大小\n",
    "        dfs(root, -1, 0)  # 0 没有父节点\n",
    "\n",
    "        def reroot(x: int, fa: int) -> None:\n",
    "            for y in g[x]:  # 遍历 x 的邻居 y\n",
    "                if y != fa:  # 避免访问父节点\n",
    "                    ans[y] = ans[x] + n - 2 * size[y]\n",
    "                    reroot(y, x)  # x 是 y 的父节点\n",
    "        reroot(root, -1)  # 0 没有父节点\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Node:\n",
    "    def __init__(self, val, parent=None):\n",
    "        self.val = val\n",
    "        self.num_descendant = 0\n",
    "        self.num_moves = 0\n",
    "        self.parent = parent\n",
    "        self.children = set()\n",
    "    \n",
    "    def __repr__(self):\n",
    "        return f\"{self.val} ({self.num_moves} children): [{self.children}]\"\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        graph = {i: set() for i in range(n)}\n",
    "        for a, b in edges:\n",
    "            graph[a].add(b)\n",
    "            graph[b].add(a)\n",
    "        \n",
    "        # Create Tree\n",
    "        queue = [Node(0)]\n",
    "        visited = {0}\n",
    "        order = [queue[0]]\n",
    "        while queue:\n",
    "            new_queue = list()\n",
    "            for node in queue:\n",
    "                for nbr in graph[node.val]:\n",
    "                    if nbr in visited:\n",
    "                        continue\n",
    "                    visited.add(nbr)\n",
    "                    new_queue.append(Node(nbr, node))\n",
    "                    node.children.add(new_queue[-1])\n",
    "            order += new_queue\n",
    "            queue = new_queue\n",
    "        order = order[::-1]\n",
    "        \n",
    "        # Calculate descendants\n",
    "        for node in order:\n",
    "            if not node.children:\n",
    "                continue\n",
    "            for child in node.children:\n",
    "                node.num_descendant += child.num_descendant + 1\n",
    "                node.num_moves += child.num_descendant + 1\n",
    "            for child in node.children:\n",
    "                node.num_moves += child.num_moves\n",
    "        \n",
    "        result = [-1] * n\n",
    "        result[0] = order[-1].num_moves\n",
    "        for i in range(len(order) - 2, -1, -1):\n",
    "            node = order[i]\n",
    "            node.num_moves += node.parent.num_moves - node.num_moves - node.num_descendant - 1 + n - node.num_descendant - 1\n",
    "            result[node.val] = node.num_moves \n",
    "\n",
    "        return result\n",
    "\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def dfs_forward(self, point, root, forwards, links):\n",
    "        forwards[point] = [1, 0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_forward(p, point, forwards, links)\n",
    "                forwards[point][0] += forwards[p][0] # num\n",
    "                forwards[point][1] += (forwards[p][0] + forwards[p][1]) # distance_sum\n",
    "        return \n",
    "    def dfs_backward(self, point, root, backwards, forwards, links):\n",
    "        # + (n-k) - k \n",
    "        backwards[point] = backwards[root] + len(backwards) - 2 * forwards[point][0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_backward(p, point, backwards, forwards, links)\n",
    "        return  \n",
    "\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if len(edges) == 0:\n",
    "            return [0]\n",
    "        # links = collections.defaultdict(list)\n",
    "        links = [[] for _ in range(n)]\n",
    "        for p1, p2 in edges:\n",
    "            links[p1].append(p2)\n",
    "            links[p2].append(p1)\n",
    "        p0 = edges[0][0]\n",
    "        forwards = [0 for i in range(n)]\n",
    "        self.dfs_forward(p0, -1, forwards, links)\n",
    "        backwards = [0 for i in range(n)]\n",
    "        backwards[p0] = forwards[p0][1]\n",
    "        for p in links[p0]:\n",
    "            self.dfs_backward(p, p0, backwards, forwards, links)\n",
    "        return backwards\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def dfs_forward(self, point, root, forwards, links):\n",
    "        forwards[point] = [1, 0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_forward(p, point, forwards, links)\n",
    "                forwards[point][0] += forwards[p][0] # num\n",
    "                forwards[point][1] += (forwards[p][0] + forwards[p][1]) \n",
    "        return \n",
    "    def dfs_backward(self, point, root, backwords, forwards, links):\n",
    "        # + (n-k) - k \n",
    "        backwords[point] = backwords[root] + len(backwords) - 2 * forwards[point][0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_backward(p, point, backwords, forwards, links)\n",
    "        return  \n",
    "\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if len(edges) == 0:\n",
    "            return [0]\n",
    "        # links = collections.defaultdict(list)\n",
    "        links = [[] for _ in range(n)]\n",
    "        for p1, p2 in edges:\n",
    "            # if p1 not in links:\n",
    "            #     links[p1] = []\n",
    "            # if p2 not in links:\n",
    "            #     links[p2] = []\n",
    "            links[p1].append(p2)\n",
    "            links[p2].append(p1)\n",
    "        p0 = edges[0][0]\n",
    "        forwards = [0 for i in range(n)]\n",
    "        self.dfs_forward(p0, -1, forwards, links)\n",
    "        backwords = [0 for i in range(n)]\n",
    "        backwords[p0] = forwards[p0][1]\n",
    "        for p in links[p0]:\n",
    "            self.dfs_backward(p, p0, backwords, forwards, links)\n",
    "        return backwords\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        '''\n",
    "        思路：递归查询，对于节点i，返回i子树的节点个数以及子树距离之和，\n",
    "             则i的所有子树距离之和为子树距离之和+子树个数\n",
    "\n",
    "             先计算出每个节点的子树距离之和，然后再递归一次，加上父亲到子的距离之和\n",
    "             父亲到自己的距离之和 = 父亲距离之和 - 该节点子树距离之和 - 该节点子树个数 + n - 该节点子树个数\n",
    "        '''\n",
    "        self.n = n\n",
    "        self.node = [[] for i in range(n)]\n",
    "        for item in edges: # 把边和节点连接起来\n",
    "            self.node[item[0]].append(item[1])\n",
    "            self.node[item[1]].append(item[0])\n",
    "        \n",
    "        self.ret = [0 for i in range(n)]\n",
    "        self.num = [0 for i in range(n)]\n",
    "        self.dfs(-1, 0) # 从根节点开始搜索\n",
    "        self.dfs2(-1, 0)\n",
    "        return self.ret\n",
    "    \n",
    "    def dfs(self, fu, u):\n",
    "        node_num = 0\n",
    "        for item in self.node[u]:\n",
    "            if item == fu:\n",
    "                continue\n",
    "            \n",
    "            sum, num = self.dfs(u, item)\n",
    "            self.ret[u] += sum\n",
    "            node_num += num\n",
    "\n",
    "        self.ret[u] += node_num\n",
    "        self.num[u] = node_num + 1\n",
    "\n",
    "        return self.ret[u], self.num[u]\n",
    "    \n",
    "    def dfs2(self, fu, u):\n",
    "        if fu != -1:\n",
    "            self.ret[u] = self.ret[fu] - self.num[u] + self.n - self.num[u]\n",
    "        for item in self.node[u]:\n",
    "            if item == fu:\n",
    "                continue\n",
    "            self.dfs2(u, item)\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def dfs_forward(self, point, root, forwards, links):\n",
    "        forwards[point] = [1, 0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_forward(p, point, forwards, links)\n",
    "                forwards[point][0] += forwards[p][0] # num\n",
    "                forwards[point][1] += (forwards[p][0] + forwards[p][1]) \n",
    "        return \n",
    "    def dfs_backward(self, point, root, backwords, forwards, links):\n",
    "        # + (n-k) - k \n",
    "        backwords[point] = backwords[root] + len(backwords) - 2 * forwards[point][0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_backward(p, point, backwords, forwards, links)\n",
    "        return  \n",
    "\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if len(edges) == 0:\n",
    "            return [0]\n",
    "        # links = collections.defaultdict(list)\n",
    "        links = dict()\n",
    "        for p1, p2 in edges:\n",
    "            if p1 not in links:\n",
    "                links[p1] = []\n",
    "            if p2 not in links:\n",
    "                links[p2] = []\n",
    "            links[p1].append(p2)\n",
    "            links[p2].append(p1)\n",
    "        p0 = edges[0][0]\n",
    "        forwards = [0 for i in range(n)]\n",
    "        self.dfs_forward(p0, -1, forwards, links)\n",
    "        backwords = [0 for i in range(n)]\n",
    "        backwords[p0] = forwards[p0][1]\n",
    "        for p in links[p0]:\n",
    "            self.dfs_backward(p, p0, backwords, forwards, links)\n",
    "        return backwords\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        g = [[] for _ in range(n)]  # g[x] 表示 x 的所有邻居\n",
    "        \n",
    "\n",
    "        # k: random root.\n",
    "        k = random.randint(0, n-1)\n",
    "\n",
    "        # 邻接表\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "\n",
    "        # for i in range(n):\n",
    "        #     if len(g[i]) > len(g[k]):\n",
    "        #         k = i\n",
    "\n",
    "        ans = [0] * n\n",
    "        size = [1] * n  # 注意这里初始化成 1 了，下面只需要累加儿子的子树大小\n",
    "        def dfs(x: int, fa: int, depth: int) -> None:\n",
    "            ans[k] += depth  # depth 为 k 到 x 的距离\n",
    "            for y in g[x]:  # 遍历 x 的邻居 y\n",
    "                if y != fa:  # 避免访问父节点\n",
    "                    dfs(y, x, depth + 1)  # x 是 y 的父节点\n",
    "                    size[x] += size[y]  # 累加 x 的儿子 y 的子树大小\n",
    "        dfs(k, -1, 0)  # 0 没有父节点\n",
    "\n",
    "        def reroot(x: int, fa: int) -> None:\n",
    "            for y in g[x]:  # 遍历 x 的邻居 y\n",
    "                if y != fa:  # 避免访问父节点\n",
    "                    ans[y] = ans[x] + n - 2 * size[y]\n",
    "                    reroot(y, x)  # x 是 y 的父节点\n",
    "        reroot(k, -1)  # 0 没有父节点\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        root = random.randint(0, n - 1)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        \n",
    "        son_cnt = [1] * n\n",
    "        dis_cnt = [0] * n\n",
    "\n",
    "        def dfs(u, fa):\n",
    "            for v in g[u]:\n",
    "                if v == fa: continue\n",
    "                dfs(v, u)\n",
    "                son_cnt[u] += son_cnt[v]\n",
    "                dis_cnt[u] += dis_cnt[v] + son_cnt[v]\n",
    "        \n",
    "        dfs(root, root)\n",
    "\n",
    "        def calc(u, fa):\n",
    "            if u != fa:\n",
    "                dis_cnt[u] = dis_cnt[fa] + (n - 2 * son_cnt[u])\n",
    "            for v in g[u]:\n",
    "                if v == fa: continue\n",
    "                calc(v, u)\n",
    "        \n",
    "        calc(root, root)\n",
    "\n",
    "        return dis_cnt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        ans = [0] * n\n",
    "        next_map = [list() for _ in range(n)]\n",
    "        for edge in edges:\n",
    "            next_map[edge[0]].append(edge[1])\n",
    "            next_map[edge[1]].append(edge[0])\n",
    "        subs = [0] * n\n",
    "\n",
    "        def func(parent, node, subs):\n",
    "            sub, val = 0, 0\n",
    "            next_list = next_map[node]\n",
    "            for next in next_list:\n",
    "                if next == parent: continue\n",
    "                s, v = func(node, next, subs)\n",
    "                sub += s + 1\n",
    "                val += v\n",
    "            subs[node] = sub\n",
    "            return sub, val + sub\n",
    "\n",
    "        ans[0] = func(-1, 0, subs)[1]\n",
    "        q = deque()\n",
    "        for next in next_map[0]:\n",
    "            q.append((next, ans[0], 0))  # idx-val-parent\n",
    "        while q:\n",
    "            idx, val, parent = q.popleft()\n",
    "            next_list = next_map[idx]\n",
    "            ans[idx] = val + (n - 1 - subs[idx] - 1) - subs[idx]\n",
    "            for next in next_list:\n",
    "                if next == parent: continue\n",
    "                q.append((next, ans[idx], idx))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        dct=[[]for i in range(n)]\n",
    "        for a,b in edges:\n",
    "            dct[a].append(b)\n",
    "            dct[b].append(a)\n",
    "        h=[-1]*n\n",
    "        c=[1]*n\n",
    "        def dfs(x):\n",
    "            for j in dct[x]:\n",
    "                if h[j]==-1:\n",
    "                    h[j]=h[x]+1\n",
    "                    c[x]+=dfs(j)\n",
    "            return c[x]\n",
    "        h[0]=0\n",
    "        dfs(0)\n",
    "        ans=[-1]*n\n",
    "        ans[0]=sum(h)\n",
    "        visit=deque()\n",
    "        visit.append(0)\n",
    "        while visit:\n",
    "            x=visit.popleft()\n",
    "            for j in dct[x]:\n",
    "                if ans[j]==-1:\n",
    "                    ans[j]=ans[x]+n-2*c[j]\n",
    "                    visit.append(j)\n",
    "                    \n",
    "        return ans\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: list) -> list:\n",
    "        from collections import defaultdict\n",
    "        dp = [0] * n\n",
    "        sz = [1] * n\n",
    "        ans = [0] * n\n",
    "        adj_dict = defaultdict(list)\n",
    "        # 预处理为邻接dict,注意是无向图，双向都要加\n",
    "        for edge in edges:\n",
    "            adj_dict[edge[0]].append(edge[1])\n",
    "            adj_dict[edge[1]].append(edge[0])\n",
    "\n",
    "        # 以root为根进行dp\n",
    "        def dfs(node: int, f_node: int):\n",
    "            if not adj_dict[node]:\n",
    "                # 叶节点dp为0, sz为1\n",
    "                return 0, 1\n",
    "            for son in adj_dict[node]:\n",
    "                # 避免访问父节点，导致死循环\n",
    "                if son == f_node:\n",
    "                    continue\n",
    "                d, s = dfs(son, node)\n",
    "                dp[node] += d + s\n",
    "                sz[node] += s\n",
    "            return dp[node], sz[node]\n",
    "\n",
    "        def dfs2(node: int, f_node: int):\n",
    "            if not adj_dict[node]:\n",
    "                return\n",
    "            # son 和 其父节点node换根\n",
    "            for son in adj_dict[node]:\n",
    "                # 避免访问父节点，导致死循环\n",
    "                if son == f_node:\n",
    "                    continue\n",
    "                dp_node, sz_node, dp_son, sz_son = dp[node], sz[node], dp[son], sz[son]\n",
    "                dp[node] -= dp[son] + sz[son]\n",
    "                sz[node] -= sz[son]\n",
    "                dp[son] += dp[node] + sz[node]\n",
    "                sz[son] += sz[node]\n",
    "                ans[son] = dp[son]\n",
    "                dfs2(son, node)\n",
    "                dp[node], sz[node], dp[son], sz[son] = dp_node, sz_node, dp_son, sz_son\n",
    "        if edges:\n",
    "            root = edges[0][0]\n",
    "            dfs(root, -1)\n",
    "            ans[root] = dp[root]\n",
    "            dfs2(root, -1)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        n = len(edges)+1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        sz = [1]*n\n",
    "        def dfs(u, pa):\n",
    "            s = 0\n",
    "            for v in g[u]:\n",
    "                if v != pa:\n",
    "                    s += dfs(v, u)+sz[v]\n",
    "                    sz[u] += sz[v]\n",
    "            return s\n",
    "        res = [0]*n\n",
    "        # res[0] = dfs(0, -1)\n",
    "        # q = deque([(0, -1)])\n",
    "        # while q:\n",
    "        #     u, pa = q.popleft()\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             res[v] = res[u]+n-2*sz[v]\n",
    "        #             q.append((v, u))\n",
    "        q = deque([(0, -1, dfs(0, -1))])\n",
    "        while q:\n",
    "            u, pa, s = q.popleft()\n",
    "            res[u] = s\n",
    "            for v in g[u]:\n",
    "                if v != pa:\n",
    "                    q.append((v, u, s+n-2*sz[v]))\n",
    "        return res\n",
    "\n",
    "        # n = len(edges)+1\n",
    "        # g = [[] for _ in range(n)]\n",
    "        # for u, v in edges:\n",
    "        #     g[u].append(v)\n",
    "        #     g[v].append(u)\n",
    "        # sz = [1]*n\n",
    "        # def dfs(u, pa):\n",
    "        #     s = 0\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             s += dfs(v, u)+sz[v]\n",
    "        #             sz[u] += sz[v]\n",
    "        #     return s\n",
    "        # res = [0]*n\n",
    "        # res[0] = dfs(0, -1)\n",
    "        # def reroot(u, pa):\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             res[v] = res[u]+n-2*sz[v]\n",
    "        #             reroot(v, u)\n",
    "        # reroot(0, -1)\n",
    "        # return res\n",
    "\n",
    "        # n = len(edges)+1\n",
    "        # g = [[] for _ in range(n)]\n",
    "        # for u, v in edges:\n",
    "        #     g[u].append(v)\n",
    "        #     g[v].append(u)\n",
    "        # sz = [1]*n\n",
    "        # def dfs(u, pa):\n",
    "        #     s = 0\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             s += dfs(v, u)+sz[v]\n",
    "        #             sz[u] += sz[v]\n",
    "        #     return s\n",
    "        # s0 = dfs(0, -1)\n",
    "        # res = [0]*n\n",
    "        # def reroot(u, pa, s):\n",
    "        #     res[u] = s\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             reroot(v, u, s+n-2*sz[v])\n",
    "        # reroot(0, -1, s0)\n",
    "        # return res\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        n = len(edges)+1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        sz = [1]*n\n",
    "        def dfs(u, pa):\n",
    "            s = 0\n",
    "            for v in g[u]:\n",
    "                if v != pa:\n",
    "                    s += dfs(v, u)+sz[v]\n",
    "                    sz[u] += sz[v]\n",
    "            return s\n",
    "        res = [0]*n\n",
    "        res[0] = dfs(0, -1)\n",
    "        q = deque([(0, -1)])\n",
    "        while q:\n",
    "            u, pa = q.popleft()\n",
    "            for v in g[u]:\n",
    "                if v != pa:\n",
    "                    res[v] = res[u]+n-2*sz[v]\n",
    "                    q.append((v, u))\n",
    "        return res\n",
    "\n",
    "        # n = len(edges)+1\n",
    "        # g = [[] for _ in range(n)]\n",
    "        # for u, v in edges:\n",
    "        #     g[u].append(v)\n",
    "        #     g[v].append(u)\n",
    "        # sz = [1]*n\n",
    "        # def dfs(u, pa):\n",
    "        #     s = 0\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             s += dfs(v, u)+sz[v]\n",
    "        #             sz[u] += sz[v]\n",
    "        #     return s\n",
    "        # res = [0]*n\n",
    "        # res[0] = dfs(0, -1)\n",
    "        # def reroot(u, pa):\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             res[v] = res[u]+n-2*sz[v]\n",
    "        #             reroot(v, u)\n",
    "        # reroot(0, -1)\n",
    "        # return res\n",
    "\n",
    "        # n = len(edges)+1\n",
    "        # g = [[] for _ in range(n)]\n",
    "        # for u, v in edges:\n",
    "        #     g[u].append(v)\n",
    "        #     g[v].append(u)\n",
    "        # sz = [1]*n\n",
    "        # def dfs(u, pa):\n",
    "        #     s = 0\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             s += dfs(v, u)+sz[v]\n",
    "        #             sz[u] += sz[v]\n",
    "        #     return s\n",
    "        # s0 = dfs(0, -1)\n",
    "        # res = [0]*n\n",
    "        # def reroot(u, pa, s):\n",
    "        #     res[u] = s\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             reroot(v, u, s+n-2*sz[v])\n",
    "        # reroot(0, -1, s0)\n",
    "        # return res\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        graph = defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            graph[a].append(b)\n",
    "            graph[b].append(a)\n",
    "\n",
    "        node_num = [1] * n\n",
    "        dist_num = [0] * n\n",
    "\n",
    "        def children_count(node, father):\n",
    "            for a in graph[node]:\n",
    "                if a == father:\n",
    "                    continue\n",
    "                children_count(a, node)\n",
    "                node_num[node] += node_num[a]\n",
    "                dist_num[node] += node_num[a] + dist_num[a]\n",
    "\n",
    "        def father_count(node, father):\n",
    "            for a in graph[node]:\n",
    "                if a == father:\n",
    "                    continue\n",
    "                dist_num[a] = dist_num[node] - node_num[a] + n - node_num[a]\n",
    "                father_count(a, node)\n",
    "\n",
    "        children_count(0, -1)\n",
    "        father_count(0, -1)\n",
    "\n",
    "        return dist_num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        child = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            child[a].append(b)\n",
    "            child[b].append(a)\n",
    "        \n",
    "        arr = [0] * n\n",
    "        def dfs(node) -> int:\n",
    "            arr[node] = -1\n",
    "            if len(lst := child[node]) == 0:\n",
    "                return 0\n",
    "            \n",
    "            ret = 0\n",
    "            for c in lst:\n",
    "                if arr[c] == 0:\n",
    "                    ret += dfs(c) - arr[c]\n",
    "                    arr[node] += arr[c]\n",
    "            \n",
    "            return ret\n",
    "        \n",
    "        arr[0] = dfs(0)\n",
    "        \n",
    "        q = deque([0])\n",
    "        while q:\n",
    "            pa = q.popleft()\n",
    "            for c in child[pa]:\n",
    "                if arr[c] < 0:\n",
    "                    arr[c] = n + arr[pa] + arr[c] * 2\n",
    "                    q.append(c)\n",
    "        \n",
    "        return arr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        n = len(edges)+1\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        sz = [1]*n\n",
    "        def dfs(u, pa):\n",
    "            s = 0\n",
    "            for v in g[u]:\n",
    "                if v != pa:\n",
    "                    s += dfs(v, u)+sz[v]\n",
    "                    sz[u] += sz[v]\n",
    "            return s\n",
    "        res = [0]*n\n",
    "        # res[0] = dfs(0, -1)\n",
    "        # q = deque([(0, -1)])\n",
    "        # while q:\n",
    "        #     u, pa = q.popleft()\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             res[v] = res[u]+n-2*sz[v]\n",
    "        #             q.append((v, u))\n",
    "        q = deque([(0, -1, dfs(0, -1))])\n",
    "        while q:\n",
    "            u, pa, s = q.popleft()\n",
    "            res[u] = s\n",
    "            for v in g[u]:\n",
    "                if v != pa:\n",
    "                    q.append((v, u, s+n-2*sz[v]))\n",
    "        return res\n",
    "\n",
    "        # n = len(edges)+1\n",
    "        # g = [[] for _ in range(n)]\n",
    "        # for u, v in edges:\n",
    "        #     g[u].append(v)\n",
    "        #     g[v].append(u)\n",
    "        # sz = [1]*n\n",
    "        # def dfs(u, pa):\n",
    "        #     s = 0\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             s += dfs(v, u)+sz[v]\n",
    "        #             sz[u] += sz[v]\n",
    "        #     return s\n",
    "        # res = [0]*n\n",
    "        # res[0] = dfs(0, -1)\n",
    "        # def reroot(u, pa):\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             res[v] = res[u]+n-2*sz[v]\n",
    "        #             reroot(v, u)\n",
    "        # reroot(0, -1)\n",
    "        # return res\n",
    "\n",
    "        # n = len(edges)+1\n",
    "        # g = [[] for _ in range(n)]\n",
    "        # for u, v in edges:\n",
    "        #     g[u].append(v)\n",
    "        #     g[v].append(u)\n",
    "        # sz = [1]*n\n",
    "        # def dfs(u, pa):\n",
    "        #     s = 0\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             s += dfs(v, u)+sz[v]\n",
    "        #             sz[u] += sz[v]\n",
    "        #     return s\n",
    "        # s0 = dfs(0, -1)\n",
    "        # res = [0]*n\n",
    "        # def reroot(u, pa, s):\n",
    "        #     res[u] = s\n",
    "        #     for v in g[u]:\n",
    "        #         if v != pa:\n",
    "        #             reroot(v, u, s+n-2*sz[v])\n",
    "        # reroot(0, -1, s0)\n",
    "        # return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        g = [[] for _ in range(n)]  # g[x] 表示 x 的所有邻居\n",
    "        \n",
    "\n",
    "        # k: random root.\n",
    "        root = random.randint(0, n-1)\n",
    "\n",
    "        # 邻接表\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "\n",
    "        # for i in range(n):\n",
    "        #     if len(g[i]) > len(g[k]):\n",
    "        #         k = i\n",
    "\n",
    "        ans = [0] * n\n",
    "        size = [1] * n  # 注意这里初始化成 1 了，下面只需要累加儿子的子树大小\n",
    "        def dfs(x: int, fa: int, depth: int) -> None:\n",
    "            ans[root] += depth  # depth 为 k 到 x 的距离\n",
    "            for y in g[x]:  # 遍历 x 的邻居 y\n",
    "                if y != fa:  # 避免访问父节点\n",
    "                    dfs(y, x, depth + 1)  # x 是 y 的父节点\n",
    "                    size[x] += size[y]  # 累加 x 的儿子 y 的子树大小\n",
    "        dfs(root, -1, 0)  # 0 没有父节点\n",
    "\n",
    "        def reroot(x: int, fa: int) -> None:\n",
    "            for y in g[x]:  # 遍历 x 的邻居 y\n",
    "                if y != fa:  # 避免访问父节点\n",
    "                    ans[y] = ans[x] + n - 2 * size[y]\n",
    "                    reroot(y, x)  # x 是 y 的父节点\n",
    "        reroot(root, -1)  # 0 没有父节点\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        graph = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            graph[a].append(b)\n",
    "            graph[b].append(a)\n",
    "        f = [0] * n\n",
    "        size = [1] * n\n",
    "        def dfs1(node, fa):\n",
    "            for x in graph[node]:\n",
    "                if x != fa:\n",
    "                    dfs1(x, node)\n",
    "                    size[node] += size[x]\n",
    "                    f[node] += size[x] + f[x]\n",
    "                    \n",
    "        \n",
    "        dfs1(0, -1)\n",
    "        def dfs2(node, fa):\n",
    "            for x in graph[node]:\n",
    "                if x != fa:\n",
    "                    f[x] = f[node] - 2 * size[x] + n\n",
    "                    dfs2(x, node)\n",
    "\n",
    "        dfs2(0, -1)\n",
    "        return f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        g=[[] for _ in range(n)]\n",
    "        for a,b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        ans=[0]*n\n",
    "        size=[1]*n\n",
    "        def dfs(x,fa,depth):\n",
    "            ans[0]+=depth\n",
    "            for y in g[x]:\n",
    "                if y!=fa:\n",
    "                    dfs(y,x,depth+1)\n",
    "                    size[x]+=size[y]\n",
    "\n",
    "        dfs(0,-1,0)\n",
    "        def reroot(x,fa):\n",
    "            for y in g[x]:\n",
    "                if y!=fa:\n",
    "                    ans[y]=ans[x]+n-2*size[y]\n",
    "                    reroot(y,x)\n",
    "        reroot(0,-1)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        # 1.DFS求每个节点的子节点\n",
    "        g = [[] for _ in range(n)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        \n",
    "        ans = [0] * n\n",
    "        size = [1] * n\n",
    "        \n",
    "        def dfs(x, fa, depth):\n",
    "            ans[0] += depth # 统计根节点至各个节点的距离\n",
    "            for y in g[x]: # 遍历邻居\n",
    "                if y != fa: # 避免访问父节点\n",
    "                    dfs(y, x, depth + 1)  # 深度递归x的子节点y\n",
    "                    size[x] += size[y]    # 当递归至叶子节点时，父节点x的子节点数量+1\n",
    "        dfs(0, -1, 0)\n",
    "        # 2.遍历每个节点，求深度，至每个子节点的数量-1，其余节点+1————ans[y] = ans[x] + n -2*子节点数量\n",
    "        def reroot(x, fa):\n",
    "            for y in g[x]:\n",
    "                if y != fa:\n",
    "                    ans[y] = ans[x] + n -2*size[y]\n",
    "                    reroot(y, x)\n",
    "        reroot(0, -1)\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        \"\"\"\n",
    "        换根dp：从0出发DFS,累积0到每个点的距离得到rs[0]\n",
    "        DFS同时计算每个子树的大小size[i]\n",
    "        再次从0来时DFS，y是x的儿子\n",
    "        那么 rs[y] = rs[x]+n-size[y]-size[y]\n",
    "        从x递推到y， 显然y到y子树内的距离都要-1 因此减去了size[y]\n",
    "        但y到y子树外的距离要增加1 y子树外有n-size[y]\n",
    "        \"\"\"\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        rs = [0 for _ in range(n)]\n",
    "        size = [1 for _ in range(n)]\n",
    "\n",
    "        def dfs(x,fa,depth):\n",
    "            rs[0] += depth #累积深度\n",
    "            for y in g[x]:\n",
    "                if y!=fa:#避免去访问父节点\n",
    "                    dfs(y,x,depth+1)\n",
    "                    size[x] += size[y]\n",
    "        \n",
    "        def dfs2(x,fa):\n",
    "             for y in g[x]:\n",
    "                if y!=fa:#避免去访问父节点\n",
    "                    rs[y] = rs[x]+n-2*size[y]\n",
    "                    dfs2(y,x)\n",
    "        \n",
    "        dfs(0,-1,0)\n",
    "        dfs2(0,-1)\n",
    "        return rs\n",
    "\n",
    "\n",
    "\n",
    "        dfs(0,-1,0)\n",
    "\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        def dfs1(parent,root,childNode,dp,size):\n",
    "            for child in childNode[root]:\n",
    "                if child != parent:\n",
    "                    dp,size = dfs1(root,child,childNode,dp,size)\n",
    "                    size[root] += size[child]\n",
    "                    dp[root] += (dp[child] + size[child])\n",
    "            return dp,size\n",
    "        \n",
    "        # def dfs2(parent,root,childNode,dp,size,res):\n",
    "        #     dp[parent] -= (dp[root] + size[root])\n",
    "        #     size[parent] -= (size[root])\n",
    "        #     dp[root] += (dp[parent] + size[parent])\n",
    "        #     size[root] += size[parent]\n",
    "        #     res[root] = dp[root]\n",
    "        #     for child in childNode[root]:\n",
    "        #         dpTmp = dp.copy()\n",
    "        #         sizeTmp = size.copy()\n",
    "        #         if child != parent:\n",
    "        #             res = dfs2(root,child,childNode,dpTmp,sizeTmp,res)\n",
    "        #     return res\n",
    "        def dfs2(parent,root,childNode,size,res):\n",
    "            res[root] = res[parent] - size[root] + n - size[root]\n",
    "            for child in childNode[root]:\n",
    "                if child != parent:\n",
    "                    dfs2(root,child,childNode,size,res)\n",
    "        \n",
    "        if n == 1:\n",
    "            return [0]\n",
    "\n",
    "        res = [-1] * n\n",
    "        dp = [0] * n\n",
    "        size = [1] * n\n",
    "        childNode = [[] for i in range(n)]\n",
    "        for edge in edges:\n",
    "            childNode[edge[0]].append(edge[1])\n",
    "            childNode[edge[1]].append(edge[0])\n",
    "        #第一遍dfs\n",
    "        for child in childNode[0]:\n",
    "            dp,size = dfs1(0,child,childNode,dp,size)\n",
    "            size[0] += size[child]\n",
    "            dp[0] += (dp[child] + size[child])\n",
    "        #换根dp-dfs\n",
    "        res[0] = dp[0]\n",
    "        for child in childNode[0]:\n",
    "            #res = dfs2(0,child,childNode,dpTmp,sizeTmp,res)\n",
    "            dfs2(0,child,childNode,size,res)\n",
    "        return res\n",
    "\n",
    "\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "\n",
    "        ad=collections.defaultdict(list)\n",
    "        for s,e in edges:\n",
    "            ad[s].append(e)\n",
    "            ad[e].append(s)\n",
    "        ans=[0]*n\n",
    "        size=[1]*n\n",
    "\n",
    "        def dfs(cur,fa,d):\n",
    "\n",
    "            ans[0]+=d \n",
    "            for nex in ad[cur]:\n",
    "                if nex != fa:\n",
    "                    dfs(nex,cur,d+1)\n",
    "                    size[cur]+=size[nex]\n",
    "        dfs(0,-1,0)\n",
    "        \n",
    "        def reroot(cur,fa):\n",
    "            for nex in ad[cur]:\n",
    "                if nex != fa:\n",
    "                    ans[nex]=ans[cur]+n-2*size[nex]\n",
    "                    reroot(nex,cur)\n",
    "        reroot(0,-1)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "\n",
    "\n",
    "        tree = [[] for _ in range(n)]\n",
    "\n",
    "        for edge in edges:\n",
    "            a, b = edge[0], edge[1]\n",
    "            tree[a].append(b) \n",
    "            tree[b].append(a) \n",
    "\n",
    "    \n",
    "        # 以0为root 求出每个子树的size\n",
    "        subtree = [1]*n\n",
    "        parents = [-1]*n\n",
    "        def dfs(root, parent):\n",
    "            for child in tree[root]:\n",
    "                if child == parent:\n",
    "                    continue\n",
    "                parents[child] = root\n",
    "                subtree[root] += dfs(child, root)\n",
    "            return subtree[root]\n",
    "        dfs(0, -1)\n",
    "        # print(tree)\n",
    "        # print()\n",
    "        # print(subtree)\n",
    "        # print(parents)\n",
    "\n",
    "        distance = [0]*n\n",
    "        def dfs2(root, path, parent):\n",
    "            res = 0\n",
    "            for i in tree[root]:\n",
    "                if i == parent:\n",
    "                    continue\n",
    "                distance[i] = path+1\n",
    "                res += dfs2(i, path+1, root)\n",
    "            return res\n",
    "        dfs2(0, 0, -1)\n",
    "\n",
    "        res = [0]*(n+1)\n",
    "        res[0] = sum(distance)\n",
    "\n",
    "        # 应该顺着parent的顺序更新\n",
    "        # for i in range(1, n):\n",
    "        #     s = subtree[i]\n",
    "        #     res[i] = res[parents[i]] - s + (n-s)\n",
    "\n",
    "        def dfs3(root, p):\n",
    "            # print(root)\n",
    "            for child in tree[root]:\n",
    "                if child == p:\n",
    "                    continue\n",
    "                s = subtree[child]\n",
    "                \n",
    "                res[child] = res[root] - s  + (n-s)\n",
    "                # print(\"-\", child, res[root], res[child])\n",
    "                dfs3(child, root)\n",
    "        # for c in tree[0]:\n",
    "        #     dfs3(c, 0)\n",
    "        dfs3(0, -1)\n",
    "        return res[:n]\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        g = defaultdict(list)\n",
    "        down = [0] * n\n",
    "        up = [0] * n\n",
    "        size = [0] * n\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        def dfs(u,fa):    \n",
    "            down[u]  = 0\n",
    "            size[u] = 0\n",
    "            for v in  g[u]:\n",
    "                if v == fa: continue\n",
    "                dfs(v,u)\n",
    "                down[u] += down[v] + size[v]\n",
    "                size[u] += size[v]\n",
    "            size[u]+= 1\n",
    "        dfs(0,-1)\n",
    "        def dfs2(u,fa):\n",
    "            for v in g[u]:\n",
    "                if v == fa: continue\n",
    "                up[v] = up[u] + down[u] - (down[v] + size[v]) + n - size[v]\n",
    "                dfs2(v,u)\n",
    "        dfs2(0,-1)        \n",
    "        ans  = []\n",
    "        for u in range(n):\n",
    "            ans.append(down[u] + up[u])\n",
    "        return ans\n",
    "\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        down = [0] * n\n",
    "        size = [1] * n\n",
    "        up = [0] * n\n",
    "        g = defaultdict(list)\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        def dfs(u,fa):\n",
    "            for v in g[u]:\n",
    "                if v != fa:\n",
    "                    dfs(v,u)\n",
    "                    down[u] += down[v] + size[v]\n",
    "                    size[u] += size[v]\n",
    "        dfs(0,-1)\n",
    "        def dfs2(u,fa):\n",
    "            for v in g[u]:\n",
    "                if v != fa:\n",
    "                    up[v] = up[u] + down[u] - (down[v] + size[v]) + n - size[v]\n",
    "                    dfs2(v,u)\n",
    "        dfs2(0,-1)\n",
    "        ans = []\n",
    "        for i in range(n):\n",
    "            ans.append(down[i] + up[i])\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        down = [0] * n\n",
    "        size = [1] * n\n",
    "        up = [0] * n\n",
    "        ans = []\n",
    "        g = defaultdict(list)\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        def dfs(u,fa):\n",
    "            for v in g[u]: \n",
    "                if v != fa:\n",
    "                    dfs(v,u)\n",
    "                    down[u] += down[v] + size[v]\n",
    "                    size[u] += size[v]\n",
    "        dfs(0,-1)\n",
    "        def dfs2(u,fa):\n",
    "            for v in g[u]:\n",
    "                if v != fa:\n",
    "                    up[v] = up[u] + down[u] - (down[v] + size[v]) + n - size[v]\n",
    "                    dfs2(v,u)\n",
    "        dfs2(0,-1)            \n",
    "\n",
    "        for i in range(n):\n",
    "            ans.append(down[i] + up[i])\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "\n",
    "        p=[[] for _ in range(n)] #图结构存储\n",
    "        g=[0]*n #g[i]来自节点i上面的路径数目\n",
    "        f=[0]*n #f[i]来自节点i下面的路径数目\n",
    "        c=[0]*n #c[i]节点i下面的所有子节点数目\n",
    "\n",
    "        #p[i]中存储着i的子节点\n",
    "        for e in edges:\n",
    "            a=e[0]\n",
    "            b=e[1]\n",
    "            p[a].append(b)\n",
    "            p[b].append(a)\n",
    "\n",
    "\n",
    "        #自底向上，f[i]的更新依赖于它下面相邻子节点的贡献个数\n",
    "        def dfs(node,fa):\n",
    "            for son in p[node]:\n",
    "                if son==fa:\n",
    "                    continue\n",
    "                dfs(son,node)\n",
    "                c[node]+=c[son]+1\n",
    "                f[node]+=f[son]+c[son]+1\n",
    "        \n",
    "        #自顶向下,g[i]的更新来自于它上面节点的g[i]和f[i]\n",
    "        def dfs1(node,fa):\n",
    "            for son in p[node]:\n",
    "                if son==fa:\n",
    "                    continue\n",
    "                l_c=n-c[son]-1\n",
    "                l_r=f[node]+g[node]-f[son]-c[son]-1\n",
    "                g[son]=l_r+l_c\n",
    "                dfs1(son,node)\n",
    "        \n",
    "        #节点i的总路数为f[i]+g[i]\n",
    "        dfs(0,-1)\n",
    "        dfs1(0,-1)\n",
    "        ans=[]\n",
    "        for i in range(n):\n",
    "            ans.append(f[i]+g[i])\n",
    "        return ans\n",
    "\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "from collections import deque, defaultdict\n",
    "from functools import lru_cache\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        \"\"\"\n",
    "\n",
    "        :param n:\n",
    "        :param edges:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        m1 = defaultdict(list)\n",
    "        for i in edges:\n",
    "            m1[i[0]].append(i[1])\n",
    "            m1[i[1]].append(i[0])\n",
    "\n",
    "        def dfs_height(node, prev, cnt):\n",
    "            cnt1 = cnt\n",
    "            for i in m1[node]:\n",
    "                if i != prev:\n",
    "                    cnt1 += dfs_height(i, node, cnt+1)\n",
    "            return cnt1\n",
    "\n",
    "        ret = []\n",
    "        ans1 = dfs_height(0, -1, 0)\n",
    "\n",
    "        mtree = defaultdict(int)\n",
    "        def iter_size(node, prev):\n",
    "            cnt1 = 1\n",
    "            for i in m1[node]:\n",
    "                if i != prev:\n",
    "                    cnt1 += iter_size(i, node)\n",
    "            mtree[node] = cnt1\n",
    "            return cnt1\n",
    "\n",
    "        iter_size(0, -1)\n",
    "        ret.append(ans1)\n",
    "        mm = {}\n",
    "        def dp_res(node, ans1, prev):\n",
    "            for i in m1[node]:\n",
    "                if i != prev:\n",
    "                    ans2 = ans1 - mtree[i] + (n - mtree[i])\n",
    "                    mm[i] = ans2\n",
    "                    dp_res(i, ans2, node)\n",
    "        dp_res(0, ans1, -1)\n",
    "        for i in range(1, n):\n",
    "            ret.append(mm[i])\n",
    "        return ret\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        graph = [set() for _ in range(n)]\n",
    "        for father, child in edges:\n",
    "            graph[father].add(child)\n",
    "            graph[child].add(father)\n",
    "        # 以0为起点构建一个存储subnodes和depth的数组\n",
    "        depth = [0] * n\n",
    "        count = [1] * n\n",
    "\n",
    "        def dfs_count_depth(node, parent=None):\n",
    "            for child in graph[node]:\n",
    "                if not child == parent:\n",
    "                    depth[child] = depth[node] + 1\n",
    "                    dfs_count_depth(child, node)\n",
    "                    count[node] += count[child]\n",
    "        dfs_count_depth(0)\n",
    "\n",
    "        ans = [0] * n\n",
    "        ans[0] = sum(depth)\n",
    "\n",
    "        def dfs_ans(node, parent=None):\n",
    "            for child in graph[node]:\n",
    "                if child != parent:\n",
    "                    ans[child] = ans[node] + n - count[child] - count[child]\n",
    "                    dfs_ans(child, node)\n",
    "        dfs_ans(0)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        G = [[] for _ in range(n)]\n",
    "        for a, b in edges:\n",
    "            G[a].append(b)\n",
    "            G[b].append(a)\n",
    "        ans, child = [0]*n, [[0, 0] for _ in range(n)]\n",
    "        def dfs1(c, p):\n",
    "            cnt, dis = 1, 0\n",
    "            for nei in G[c]:\n",
    "                if nei!=p:\n",
    "                    dfs1(nei, c)\n",
    "                    cnt += child[nei][0]\n",
    "                    dis += sum(child[nei])\n",
    "            child[c] = [cnt, dis]\n",
    "        def dfs2(c, p, pdis):\n",
    "            print(c, p, pdis)\n",
    "            ans[c] += child[c][1]\n",
    "            pdis += n-child[c][0]\n",
    "            ans[c] += pdis \n",
    "            for nei in G[c]:\n",
    "                if nei!=p:\n",
    "                    dfs2(nei, c, pdis+child[c][1]-child[nei][1]-child[nei][0])\n",
    "        root = 0\n",
    "        dfs1(root, -1)\n",
    "        dfs2(root, -1, 0)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        child = [1] * n\n",
    "        ans = [0] * n\n",
    "        tree = [set() for _ in range(n)]\n",
    "\n",
    "        for a, b in edges:\n",
    "            tree[a].add(b)\n",
    "            tree[b].add(a)\n",
    "        \n",
    "        def dfs(node, father, depth):\n",
    "            ans[0] += depth\n",
    "            if father in tree[node]:\n",
    "                tree[node].remove(father)\n",
    "\n",
    "            for a in tree[node]:\n",
    "                child[node] += dfs(a, node, depth+1)\n",
    "            \n",
    "            return child[node]\n",
    "        \n",
    "        dfs(0, -1, 0)\n",
    "\n",
    "        def dfs1(node, father):\n",
    "            ans[node] = ans[father] + n - 2 * child[node]\n",
    "\n",
    "            for a in tree[node]:\n",
    "                dfs1(a, node)\n",
    "        \n",
    "        for a in tree[0]:\n",
    "            dfs1(a, 0)\n",
    "        \n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        g = defaultdict(set)\n",
    "        len = n\n",
    "        for a, b in edges:\n",
    "            g[a].add(b)\n",
    "            g[b].add(a)\n",
    "\n",
    "        dp = [0] * n\n",
    "        sz = [0] * n\n",
    "        def init(node, pre, depth):\n",
    "            dp[0] += depth\n",
    "\n",
    "            sz[node] = 1\n",
    "            for n in g[node]:\n",
    "                if n == pre: continue\n",
    "                init(n, node, depth + 1)\n",
    "                sz[node] += sz[n]\n",
    "\n",
    "        init(0, -1, 0)\n",
    "\n",
    "        def dfs(node, pre):            \n",
    "            for n in g[node]:\n",
    "                if n == pre: continue\n",
    "                dp[n] = dp[node] + (len - sz[n]) - sz[n]\n",
    "                dfs(n, node)\n",
    "\n",
    "        dfs(0, -1)\n",
    "        return dp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        graph = defaultdict(list)\n",
    "        for a, b in edges:\n",
    "            graph[a].append(b)\n",
    "            graph[b].append(a)\n",
    "        subtree_sizes = [0 for _ in range(n)]\n",
    "        res = [0 for _ in range(n)]\n",
    "        visited = set()\n",
    "        def get_subtree_size(node, dist):\n",
    "            visited.add(node)\n",
    "            nonlocal res\n",
    "            nonlocal subtree_sizes\n",
    "            res[0] += dist\n",
    "            if node not in graph:\n",
    "                return 1\n",
    "            size = 1\n",
    "            for nei in graph[node]:\n",
    "                if nei not in visited:\n",
    "                    size += get_subtree_size(nei, dist+1)\n",
    "            subtree_sizes[node] = size\n",
    "            return size\n",
    "        get_subtree_size(0, 0)\n",
    "        visited = set([0])\n",
    "        # print(subtree_sizes, res)\n",
    "        def get_res(node, parent):\n",
    "            visited.add(node)\n",
    "            res[node] = res[parent] + n - 2 * subtree_sizes[node]\n",
    "            for nei in graph[node]:\n",
    "                if nei not in visited:\n",
    "                    visited.add(nei)\n",
    "                    get_res(nei, node)\n",
    "        for nei in graph[0]:\n",
    "            get_res(nei, 0)\n",
    "        return res\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        graph = defaultdict(set)\n",
    "        for u, v in edges:\n",
    "            graph[u].add(v)\n",
    "            graph[v].add(u)\n",
    "        \n",
    "        size = [0] * n\n",
    "        res = [0] * n\n",
    "        def dfs(cur, parent, depth):\n",
    "            res[0] += depth\n",
    "            if not graph[cur]:\n",
    "                size[cur] = 1 \n",
    "                return 1\n",
    "            \n",
    "            cnt = 1\n",
    "            for child in graph[cur]:\n",
    "                if child != parent:\n",
    "                    cnt += dfs(child, cur, depth+1)\n",
    "            size[cur] = cnt\n",
    "            return cnt\n",
    "        \n",
    "        dfs(0, -1, 0)\n",
    "        print(size, res)\n",
    "\n",
    "        def reroot(cur, parent):\n",
    "            for child in graph[cur]:\n",
    "                if child != parent:\n",
    "                    res[child] = res[cur] + n - 2 * size[child]\n",
    "                    reroot(child, cur)\n",
    "        \n",
    "        reroot(0, -1)\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "\n",
    "        if n == 1: return [0]\n",
    "\n",
    "        w = defaultdict(dict)\n",
    "        for s,t in edges:\n",
    "            w[s][t] = 0\n",
    "            w[t][s] = 0\n",
    "        \n",
    "        def dfs(node, f):\n",
    "\n",
    "            c, v = 0, 0\n",
    "            for nxt in w[node]:\n",
    "                if nxt == f: continue\n",
    "                cc, cv = dfs(nxt, node)\n",
    "                w[node][nxt] = cc\n",
    "                c += cc\n",
    "                v += cv\n",
    "            v += c\n",
    "            c += 1\n",
    "            return c, v\n",
    "        \n",
    "        _, v = dfs(0, -1)\n",
    "        ans = [0]*n\n",
    "        ans[0] = v\n",
    "\n",
    "        def getv(node, f, v):\n",
    "            for nxt in w[node]:\n",
    "                if nxt == f: continue\n",
    "                l, r = w[node][nxt], n-1-w[node][nxt]\n",
    "                ans[nxt] = v-l+r+1\n",
    "                getv(nxt, node, ans[nxt])\n",
    "        \n",
    "        getv(0, -1, ans[0])\n",
    "        return ans\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:\n",
    "        def traverse(u):\n",
    "            visited.add(u)\n",
    "            for v in neighbors[u]:\n",
    "                if v not in visited:\n",
    "                    traverse(v)\n",
    "                    cnt[u] += cnt[v]\n",
    "                    dis[u] += dis[v] + cnt[v]\n",
    "        def dist(u):\n",
    "            visited.add(u)\n",
    "            for v in neighbors[u]:\n",
    "                if v not in visited:\n",
    "                    dis[v] = dis[u] - cnt[v] + N - cnt[v]\n",
    "                    dist(v)\n",
    "        \n",
    "        cnt, dis = [1] * N, [0] * N\n",
    "        visited = set()\n",
    "        neighbors = defaultdict(set)\n",
    "        for u,v in edges:\n",
    "            neighbors[u].add(v)\n",
    "            neighbors[v].add(u)\n",
    "        traverse(0)\n",
    "        visited.clear()\n",
    "        dist(0)\n",
    "        return dis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        up = [[0, 1] for _ in range(n)]  # (sum, sz)\n",
    "        down = [[0, 1] for _ in range(n)]  # (sum, sz)\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        def dfs1(now, pre):\n",
    "            for nei in g[now]:\n",
    "                if nei == pre:\n",
    "                    continue\n",
    "                dfs1(nei, now)\n",
    "                s, sz = down[nei]\n",
    "                down[now][0] += s + sz\n",
    "                down[now][1] += sz\n",
    "        def dfs2(now, pre):\n",
    "            for nei in g[now]:\n",
    "                if nei == pre:\n",
    "                    continue\n",
    "                s1, sz1 = up[now]\n",
    "                s2, sz2 = down[now]\n",
    "                s3, sz3 = down[nei]\n",
    "                up[nei][1] = sz1 + sz2 - sz3\n",
    "                up[nei][0] = s1 + s2 - (s3 + sz3) + up[nei][1] - 1\n",
    "                dfs2(nei, now)\n",
    "        dfs1(0, -1)\n",
    "        dfs2(0, -1)\n",
    "\n",
    "        return [u[0] + d[0] for u, d in zip(up, down)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def cal(self,base,root,tree):\n",
    "        '''计算当前节点root为根节点到真根节点的距离和,base为root和真根的距离'''\n",
    "        if len(tree[root][0])==0: #没孩子\n",
    "            return 0\n",
    "        ret = (base+1)*len(tree[root][0])\n",
    "        for c in tree[root][0]:\n",
    "            ret += self.cal(base+1,c,tree)\n",
    "        return ret\n",
    "\n",
    "    def cal2(self,cur,ret,tree):\n",
    "        parent = tree[cur][1]\n",
    "        parentV = ret[parent]\n",
    "        if parentV==None:\n",
    "            parentV = self.cal2(parent,ret,tree)\n",
    "        return parentV-2+len(ret)-2*(tree[cur][2])\n",
    "\n",
    "    def cal3(self,root,tree):\n",
    "        '''计算root的子孙数'''\n",
    "        if tree[root][2]!=None:\n",
    "            return tree[root][2]\n",
    "        tree[root][2] = len(tree[root][0])\n",
    "        for c in tree[root][0]:\n",
    "            tree[root][2] += self.cal3(c,tree)\n",
    "        return tree[root][2]\n",
    "\n",
    "    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:\n",
    "        if N==1:\n",
    "            return [0]\n",
    "        if N==2:\n",
    "            return [1,1]\n",
    "\n",
    "        ret = [None]*N\n",
    "        tree = []\n",
    "        for _ in range(N):\n",
    "            tree.append([[],None,None]) #chilren, parent, 子孙数目\n",
    "        d = [None]*N\n",
    "        for e in edges:\n",
    "            p1 = e[0]\n",
    "            p2 = e[1]\n",
    "            if d[p1]==None:\n",
    "                d[p1] = {p2}\n",
    "            else:\n",
    "                d[p1].add(p2)\n",
    "            if d[p2]==None:\n",
    "                d[p2] = {p1}\n",
    "            else:\n",
    "                d[p2].add(p1)\n",
    "        #print(d)\n",
    "        cur = {0}\n",
    "        while len(cur)!=0:\n",
    "            ncur = set()\n",
    "            for c in cur:\n",
    "                for child in d[c]:\n",
    "                    if child==tree[c][1]:\n",
    "                        pass\n",
    "                    else:\n",
    "                        tree[c][0].append(child)\n",
    "                        tree[child][1] = c\n",
    "                        ncur.add(child)\n",
    "            cur = ncur\n",
    "        #print(tree)\n",
    "        self.cal3(0,tree)\n",
    "        #print(tree)\n",
    "        ret[0] = self.cal(0,0,tree)\n",
    "        for i in range(1,N):\n",
    "            ret[i] = self.cal2(i,ret,tree)\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Node:\n",
    "    def __init__(self,val=0) -> None:\n",
    "        self.val=val\n",
    "        self.childs=set()\n",
    "        self.dis=0\n",
    "        self.NumOfChild=0\n",
    "class Solution:\n",
    "    tree=[]\n",
    "    n=0\n",
    "    def downToUp(self,node:int,parent):\n",
    "\n",
    "        if parent in self.tree[node].childs:\n",
    "                self.tree[node].childs.remove(parent)\n",
    "\n",
    "        for child in self.tree[node].childs:\n",
    "\n",
    "            getReturn=self.downToUp(child,node)\n",
    "            self.tree[node].NumOfChild+=1+getReturn[0]\n",
    "            self.tree[node].dis+=getReturn[1]+getReturn[0]+1\n",
    "        return [self.tree[node].NumOfChild,self.tree[node].dis]\n",
    "            \n",
    "    def upToDown(self,node:int,dis:int):\n",
    "        partB=self.tree[node].dis+self.tree[node].NumOfChild+1\n",
    "        another=dis-partB\n",
    "        self.tree[node].dis+=another+(self.n-self.tree[node].NumOfChild-1)\n",
    "        for child in self.tree[node].childs:\n",
    "            self.upToDown(child,self.tree[node].dis)\n",
    "\n",
    "    def sumOfDistancesInTree(self, n: int, edges: list[list[int]]) -> list[int]:\n",
    "        self.n=n\n",
    "        self.tree=[0]*n\n",
    "        ans=[]\n",
    "        for k in range(n):\n",
    "            self.tree[k]=Node(k)\n",
    "\n",
    "        for edge in edges:\n",
    "            ori,tar=edge[0],edge[1]\n",
    "            self.tree[ori].childs.add(tar)\n",
    "            self.tree[tar].childs.add(ori)\n",
    "        self.downToUp(0,-1)\n",
    "        self.upToDown(0,self.tree[0].dis+self.tree[0].NumOfChild+1)\n",
    "        for k in range(n):\n",
    "            ans.append(self.tree[k].dis)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        self.n = n\n",
    "        self.tree_size = {i:0 for i in range(n)}\n",
    "        self.graph = {i:set([]) for i in range(n)}\n",
    "        for edge in edges:\n",
    "            x,y = edge\n",
    "            self.graph[x].add(y)\n",
    "            self.graph[y].add(x)\n",
    "        self.ans = [0]*n\n",
    "        self.dfs(0, set([0]), 1)\n",
    "        # tree dp\n",
    "        self.tree_dp(0, set([0]))\n",
    "        return self.ans\n",
    "\n",
    "\n",
    "    \n",
    "    def dfs(self, root, visited, level):\n",
    "        ans = 1\n",
    "        for next_node in self.graph[root]:\n",
    "            if next_node not in visited:\n",
    "                visited.add(next_node)\n",
    "                self.ans[0] += level\n",
    "                ans += self.dfs(next_node, visited, level+1)\n",
    "        self.tree_size[root] = ans\n",
    "        return ans\n",
    "    \n",
    "    def tree_dp(self, root, visited):\n",
    "        for next_node in self.graph[root]:\n",
    "            if next_node not in visited:\n",
    "                # sum_subtree = sum dist(root, node_j) for every subtree node j based on the root\n",
    "                # sum_other = sum dist(root, node_i) for every node i exclude from the subtree base on the root\n",
    "                # sum_distance = sum_subtree + sum_other\n",
    "                # let x be the parent of y\n",
    "                # sum_distance(y) = sum_subtree_y + sum_other_y\n",
    "                # sum_subtree_y = sum_subtree_x - size_y\n",
    "                # sum_other_y = sum_other_x + n - size_y(all excluded nodes dist + 1)\n",
    "                self.ans[next_node] = self.ans[root] + self.n - 2*self.tree_size[next_node]\n",
    "                visited.add(next_node)\n",
    "                self.tree_dp(next_node, visited)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        def f(node):\n",
    "            for child in children[node]:\n",
    "                f(child)\n",
    "\n",
    "        down_counts = [1] * n  # 当前节点下方节点的个数(包括自己)\n",
    "        up_counts = [0] * n  # 当前节点上方节点的个数(包括父节点的其他子树节点，不包括自己)\n",
    "        depths = [0] * n  # 深度。0节点深度为0\n",
    "        distances = [0] * n  # 距离和\n",
    "\n",
    "        children = defaultdict(set)\n",
    "        for f, t in edges:\n",
    "            children[f].add(t)\n",
    "            children[t].add(f)\n",
    "\n",
    "        def clear_fathers(node):\n",
    "            for child in children[node]:\n",
    "                children[child].discard(node)\n",
    "                clear_fathers(child)\n",
    "\n",
    "        def cal_down_counts(node):\n",
    "            cnt = 1\n",
    "            for child in children[node]:\n",
    "                cnt += cal_down_counts(child)\n",
    "            down_counts[node] = cnt\n",
    "            return cnt\n",
    "\n",
    "        def cal_up_counts(node):\n",
    "            up_count = up_counts[node]\n",
    "            down_count = down_counts[node]\n",
    "            for child in children[node]:\n",
    "                new_up_node_cnt = down_count - down_counts[child]\n",
    "                up_counts[child] = up_count + new_up_node_cnt\n",
    "                cal_up_counts(child)\n",
    "\n",
    "        def cal_depths(node, depth):\n",
    "            depths[node] = depth\n",
    "            for child in children[node]:\n",
    "                cal_depths(child, depth + 1)\n",
    "\n",
    "        def cal_distances(node):\n",
    "            dist = distances[node]\n",
    "            for child in children[node]:\n",
    "                delta = up_counts[child] - down_counts[child]\n",
    "                distances[child] = dist + delta\n",
    "                cal_distances(child)\n",
    "\n",
    "        clear_fathers(0)\n",
    "        cal_down_counts(0)\n",
    "        cal_up_counts(0)\n",
    "        cal_depths(0, 0)\n",
    "        distances[0] = sum(depths)\n",
    "        cal_distances(0)\n",
    "\n",
    "        return distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "#\n",
    "# @lc app=leetcode.cn id=834 lang=python3\n",
    "#\n",
    "# [834] 树中距离之和\n",
    "#\n",
    "\n",
    "# @lc code=start\n",
    "class TreeNode:\n",
    "    def __init__(self) -> None:\n",
    "        self.children=[]\n",
    "        self.father = None\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges):\n",
    "        def dfs(index):\n",
    "            node = nod[index]\n",
    "            count[index] = 1\n",
    "            for i in node.children:\n",
    "                dfs(i)\n",
    "                route[index]+=route[i]+count[i]\n",
    "                count[index]+=count[i]\n",
    "        def resdfs(index):\n",
    "            node = nod[index]\n",
    "            for i in node.children:\n",
    "                res[i] = res[index]+n-2*count[i]\n",
    "                resdfs(i)\n",
    "        nod=[TreeNode()for i in range(n)]\n",
    "        dic = defaultdict(list)\n",
    "        for i in edges:\n",
    "            dic[i[0]].append(i[1])\n",
    "            dic[i[1]].append(i[0])\n",
    "        queue = deque([0])\n",
    "        vis = set()\n",
    "        vis.add(0)\n",
    "        while queue:\n",
    "            x = queue.popleft()\n",
    "            for i in dic[x]:\n",
    "                if i not in vis:\n",
    "                    nod[x].children.append(i)\n",
    "                    queue.append(i)\n",
    "                    vis.add(i)\n",
    "        route = [0]*n\n",
    "        count = [1]*n\n",
    "        dfs(0)\n",
    "        res = [0]*n\n",
    "        res[0]=route[0]\n",
    "        resdfs(0)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import deque\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        fas = [None] * n\n",
    "        chs = [[] for _ in range(n)]\n",
    "\n",
    "        edge_maps = [[] for _ in range(n)]\n",
    "\n",
    "        for edge in edges:\n",
    "            a, b = edge\n",
    "            edge_maps[a].append(b)\n",
    "            edge_maps[b].append(a)\n",
    "\n",
    "        stk = deque([0])\n",
    "        while len(stk)>0:\n",
    "            node = stk.popleft()\n",
    "            for ch in edge_maps[node]:\n",
    "                if ch == fas[node]:\n",
    "                    continue\n",
    "                chs[node].append(ch)\n",
    "                fas[ch] = node\n",
    "                stk.append(ch)\n",
    "\n",
    "        \n",
    "        hs = [[] for _ in range(n)]\n",
    "        \n",
    "        def helper(node:int) -> Tuple[int, int]:\n",
    "            if len(hs[node]) == 0:\n",
    "                ret_n, ret_sumdis = 1, 0\n",
    "                for child in chs[node]:\n",
    "                    n_chs, sumdis = helper(child)\n",
    "                    ret_n += n_chs\n",
    "                    ret_sumdis += sumdis + (n_chs-1) + 1\n",
    "\n",
    "                hs[node] = [ret_n, ret_sumdis]\n",
    "\n",
    "                return ret_n, ret_sumdis\n",
    "\n",
    "            else:\n",
    "                return hs[node]\n",
    "        \n",
    "        helper(0)\n",
    "\n",
    "        # print(hs)\n",
    "        # print('===')\n",
    "\n",
    "        return_list = [None] * n\n",
    "\n",
    "        seen = set()\n",
    "        def helper2(node: int) -> None:\n",
    "            return_list[node] = hs[node][1]\n",
    "            seen.add(node)\n",
    "            children = copy.deepcopy(chs[node])\n",
    "            # print(f'{node},{children}')\n",
    "            for ch_node in children:\n",
    "                if ch_node in seen:\n",
    "                    continue\n",
    "\n",
    "                fas[ch_node] = None\n",
    "                fas[node] = ch_node\n",
    "                chs[ch_node].append(node)\n",
    "                chs[node].remove(ch_node)\n",
    "                hs[node][0] = hs[node][0] - hs[ch_node][0]\n",
    "                hs[node][1] = hs[node][1] - (hs[ch_node][0] + hs[ch_node][1])\n",
    "                hs[ch_node][0] = hs[ch_node][0] + hs[node][0]\n",
    "                hs[ch_node][1] = hs[ch_node][1] + (hs[node][0] + hs[node][1])\n",
    "                \n",
    "                helper2(ch_node)\n",
    "\n",
    "                fas[node] = None\n",
    "                fas[ch_node] = node\n",
    "                chs[node].append(ch_node)\n",
    "                chs[ch_node].remove(node)\n",
    "                hs[ch_node][0] = hs[ch_node][0] - hs[node][0]\n",
    "                hs[ch_node][1] = hs[ch_node][1] - (hs[node][0] + hs[node][1])\n",
    "                hs[node][0] = hs[node][0] + hs[ch_node][0]\n",
    "                hs[node][1] = hs[node][1] + (hs[ch_node][0] + hs[ch_node][1])\n",
    "\n",
    "        helper2(0)\n",
    "\n",
    "\n",
    "        return return_list\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\"\"\"\n",
    "树形换根dp\n",
    "\"\"\"\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        tree = [[] for _ in range(n)]\n",
    "        sz = [1 for _ in range(n)]\n",
    "        dp = [0 for _ in range(n)]\n",
    "\n",
    "        # 树可以任一节点当根\n",
    "        # 需要存为无向图\n",
    "        for e in edges:\n",
    "            [u, v] = e\n",
    "            tree[u].append(v)\n",
    "            tree[v].append(u)\n",
    "\n",
    "        # 自底向上dp\n",
    "        def bottomToUpDP(u, fa):\n",
    "            for v in tree[u]:\n",
    "                if v == fa: continue\n",
    "                bottomToUpDP(v, u)\n",
    "                dp[u] += dp[v] + sz[v]\n",
    "                sz[u] += sz[v]\n",
    "        \n",
    "        bottomToUpDP(0, -1)\n",
    "        ret = [0 for _ in range(n)]\n",
    "\n",
    "        # 换根\n",
    "        def exchangeRootDP(u, fa):\n",
    "            ret[u] = dp[u]\n",
    "            for v in tree[u]:\n",
    "                if v == fa: continue\n",
    "                uDp, uSz = dp[u], sz[u]\n",
    "                vDp, vSz = dp[v], sz[v]\n",
    "                dp[u] -= dp[v] + sz[v]\n",
    "                sz[u] -= sz[v]\n",
    "                dp[v] += dp[u] + sz[u]\n",
    "                sz[v] += sz[u]\n",
    "                exchangeRootDP(v, u)\n",
    "                dp[u], sz[u] = uDp, uSz\n",
    "                dp[v], sz[v] = vDp, vSz\n",
    "\n",
    "        exchangeRootDP(0, -1)\n",
    "        return ret\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        ans = [0] * n\n",
    "        dp = [0] * n\n",
    "        snode = [1] * n\n",
    "        graph = [[] for _ in range(n)]\n",
    "        for u,v in edges:\n",
    "            graph[u].append(v)\n",
    "            graph[v].append(u)\n",
    "\n",
    "        def dfs(u, f):\n",
    "            \"\"\"\n",
    "            u : 当前节点\n",
    "            v ：当前节点的父节点\n",
    "            \"\"\"\n",
    "            for v in graph[u]:\n",
    "                if v != f:\n",
    "                    dfs(v, u)\n",
    "                    dp[u] += dp[v] + snode[v]\n",
    "                    snode[u] += snode[v]\n",
    "        # 从第一个节点开始寻找\n",
    "        dfs(0,-1)\n",
    "\n",
    "        def dfs2(u, f):\n",
    "            \"\"\"\n",
    "            \"\"\"\n",
    "            ans[u] = dp[u]\n",
    "            for v in graph[u]:\n",
    "                if v != f:\n",
    "                    dp_u, dp_v, snode_u, snode_v = dp[u],dp[v],snode[u],snode[v]\n",
    "                    dp[u] -= dp[v] + snode[v]\n",
    "                    snode[u] -= snode[v]\n",
    "                    dp[v] += dp[u] + snode[u]\n",
    "                    snode[v] += snode[u]\n",
    "\n",
    "                    dfs2(v,u)\n",
    "\n",
    "                    dp[u],dp[v],snode[u],snode[v] = dp_u, dp_v, snode_u, snode_v\n",
    "        dfs2(0,-1)\n",
    "        return ans"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
